Face Fit or Restore implementation
This commit is contained in:
Rafał Tłuchowski
+21
@@ -0,0 +1,21 @@
|
||||
# __init__.py
|
||||
import os
|
||||
import sys
|
||||
from .nodes.face_fit_and_restore import FaceFitAndRestore
|
||||
|
||||
# Add current directory to system path
|
||||
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
|
||||
|
||||
# Get the path to the current directory
|
||||
NODE_PATH = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FaceFitAndRestore": FaceFitAndRestore
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FaceFitAndRestore": "Face Fit or Restore"
|
||||
}
|
||||
|
||||
__version__ = "1.0.0"
|
||||
@@ -0,0 +1,66 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Set, Optional
|
||||
|
||||
|
||||
def generate_context_file(
|
||||
root_dir: str,
|
||||
output_file: str,
|
||||
ignore_dirs: Optional[Set[str]] = None,
|
||||
ignore_files: Optional[Set[str]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Generate a single file containing all code with file paths as headers.
|
||||
|
||||
Args:
|
||||
root_dir: Root directory to start scanning.
|
||||
output_file: Output file path.
|
||||
ignore_dirs: Set of directory names to ignore.
|
||||
ignore_files: Set of file names to ignore.
|
||||
"""
|
||||
# Default ignored directories and files
|
||||
if ignore_dirs is None:
|
||||
ignore_dirs = {".git", "__pycache__", ".venv", "venv", ".idea", "node_modules"}
|
||||
if ignore_files is None:
|
||||
ignore_files = set()
|
||||
|
||||
# Ensure the output directory exists
|
||||
output_path = Path(output_file)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(output_file, "w", encoding="utf-8") as out:
|
||||
for root, dirs, files in os.walk(root_dir):
|
||||
# Filter out ignored directories
|
||||
dirs[:] = [d for d in dirs if d not in ignore_dirs]
|
||||
|
||||
# Process each file
|
||||
for file in files:
|
||||
if file in ignore_files:
|
||||
continue
|
||||
|
||||
# Check if the file has a supported extension
|
||||
if file.endswith((".py", ".sql", ".js", ".tsx", ".ts", ".jsx", ".json", ".css", ".html")):
|
||||
file_path = Path(root) / file
|
||||
relative_path = file_path.relative_to(root_dir)
|
||||
|
||||
# Write file path as a header
|
||||
out.write(f"\n# {relative_path}\n")
|
||||
|
||||
# Write file content
|
||||
try:
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
out.write(f.read())
|
||||
out.write("\n")
|
||||
except Exception as e:
|
||||
out.write(f"# Error reading file: {str(e)}\n")
|
||||
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
output_file = "D:/code/code_base/face_processor_02.txt"
|
||||
generate_context_file(
|
||||
root_dir="E:/ComfyUI-DEV/custom_nodes/ComfyUI_FaceProcessor",
|
||||
output_file=output_file,
|
||||
ignore_dirs={".git", "__pycache__", ".venv", "venv", ".idea", "node_modules", "models"},
|
||||
ignore_files={".gitignore", "contex_generator.py", output_file, "package-lock.json", "TODO.txt"},
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
import mediapipe as mp
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import cv2
|
||||
import torch
|
||||
from PIL import Image
|
||||
from typing import Union, Optional
|
||||
|
||||
class FaceDetector:
|
||||
def __init__(self):
|
||||
"""Initialize the FaceDetector with MediaPipe Face Mesh."""
|
||||
self.mp_face_mesh = mp.solutions.face_mesh
|
||||
self.face_mesh = self.mp_face_mesh.FaceMesh(
|
||||
static_image_mode=True,
|
||||
max_num_faces=1,
|
||||
refine_landmarks=True,
|
||||
min_detection_confidence=0.5
|
||||
)
|
||||
|
||||
def detect_landmarks(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[pd.DataFrame]:
|
||||
"""Detect facial landmarks in the given image."""
|
||||
image_np = self._convert_to_numpy(image)
|
||||
if image_np is None:
|
||||
return None
|
||||
|
||||
results = self.face_mesh.process(image_np)
|
||||
if not results.multi_face_landmarks:
|
||||
print("No face detected in the image")
|
||||
return None
|
||||
|
||||
face_landmarks = results.multi_face_landmarks[0]
|
||||
landmarks_data = {'x': [], 'y': [], 'z': [], 'index': []}
|
||||
image_height, image_width = image_np.shape[:2]
|
||||
|
||||
for idx, landmark in enumerate(face_landmarks.landmark):
|
||||
x = landmark.x * image_width
|
||||
y = landmark.y * image_height
|
||||
z = landmark.z
|
||||
|
||||
landmarks_data['x'].append(x)
|
||||
landmarks_data['y'].append(y)
|
||||
landmarks_data['z'].append(z)
|
||||
landmarks_data['index'].append(idx)
|
||||
|
||||
return pd.DataFrame(landmarks_data)
|
||||
|
||||
def _convert_to_numpy(self, image: Union[torch.Tensor, np.ndarray, Image.Image]) -> Optional[np.ndarray]:
|
||||
"""Convert different image types to numpy array."""
|
||||
if torch.is_tensor(image):
|
||||
image = image.detach().cpu().numpy()
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
image = (image * 255).astype(np.uint8)
|
||||
elif isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif isinstance(image, np.ndarray):
|
||||
if image.dtype == np.float32 and image.max() <= 1.0:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
|
||||
if len(image.shape) == 2:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
||||
elif image.shape[2] == 4:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
|
||||
|
||||
return image
|
||||
|
||||
def __del__(self):
|
||||
"""Clean up resources."""
|
||||
if hasattr(self, 'face_mesh'):
|
||||
self.face_mesh.close()
|
||||
@@ -0,0 +1,170 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from PIL import Image
|
||||
from typing import Union, Optional, Tuple
|
||||
|
||||
class ImageProcessor:
|
||||
@staticmethod
|
||||
def _convert_to_numpy(image: Union[torch.Tensor, np.ndarray, Image.Image]) -> np.ndarray:
|
||||
"""Convert different image types to numpy array."""
|
||||
if torch.is_tensor(image):
|
||||
image = image.detach().cpu().numpy()
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
image = (image * 255).astype(np.uint8)
|
||||
elif isinstance(image, Image.Image):
|
||||
image = np.array(image)
|
||||
elif isinstance(image, np.ndarray):
|
||||
if image.dtype == np.float32 and image.max() <= 1.0:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
|
||||
if len(image.shape) == 2:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
||||
elif image.shape[2] == 4:
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
|
||||
|
||||
return image
|
||||
|
||||
@staticmethod
|
||||
def calculate_face_bbox(landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Optional[Tuple[int, int, int, int]]:
|
||||
"""Calculate bounding box for face based on landmarks."""
|
||||
if landmarks_df is None or landmarks_df.empty:
|
||||
return None
|
||||
|
||||
min_x = landmarks_df['x'].min()
|
||||
max_x = landmarks_df['x'].max()
|
||||
min_y = landmarks_df['y'].min()
|
||||
max_y = landmarks_df['y'].max()
|
||||
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
|
||||
pad_x = width * padding_percent
|
||||
pad_y = height * padding_percent
|
||||
|
||||
x1 = max(0, min_x - pad_x)
|
||||
y1 = max(0, min_y - pad_y)
|
||||
x2 = max_x + pad_x
|
||||
y2 = max_y + pad_y
|
||||
|
||||
return (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
|
||||
|
||||
@staticmethod
|
||||
def crop_face(image: Union[torch.Tensor, np.ndarray, Image.Image], landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Optional[np.ndarray]:
|
||||
"""Crop face region from image based on landmarks."""
|
||||
image_np = ImageProcessor._convert_to_numpy(image)
|
||||
bbox = ImageProcessor.calculate_face_bbox(landmarks_df, padding_percent)
|
||||
|
||||
if bbox is None:
|
||||
return None
|
||||
|
||||
x, y, w, h = bbox
|
||||
return image_np[y:y + h, x:x + w]
|
||||
|
||||
@staticmethod
|
||||
def resize_image(image: Union[torch.Tensor, np.ndarray, Image.Image], target_size: int) -> Optional[np.ndarray]:
|
||||
"""Resize image to target size while maintaining aspect ratio and cropping to square."""
|
||||
if image is None:
|
||||
return None
|
||||
|
||||
image_np = ImageProcessor._convert_to_numpy(image)
|
||||
h, w = image_np.shape[:2]
|
||||
|
||||
# Determine the smaller dimension (height or width)
|
||||
min_dim = min(h, w)
|
||||
|
||||
# Calculate the crop box to make the image square
|
||||
start_x = (w - min_dim) // 2
|
||||
start_y = (h - min_dim) // 2
|
||||
end_x = start_x + min_dim
|
||||
end_y = start_y + min_dim
|
||||
|
||||
# Crop the image to a square
|
||||
cropped_image = image_np[start_y:end_y, start_x:end_x]
|
||||
|
||||
# Resize the cropped image to the target size
|
||||
resized = cv2.resize(cropped_image, (target_size, target_size), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
return resized
|
||||
|
||||
@staticmethod
|
||||
def calculate_rotation_angle(landmarks_df: pd.DataFrame) -> float:
|
||||
"""Calculate rotation angle based on eyes position."""
|
||||
if landmarks_df is None or landmarks_df.empty:
|
||||
return 0.0
|
||||
|
||||
LEFT_EYE = 33 # Center of the left eye
|
||||
RIGHT_EYE = 263 # Center of the right eye
|
||||
|
||||
left_eye = landmarks_df[landmarks_df['index'] == LEFT_EYE].iloc[0]
|
||||
right_eye = landmarks_df[landmarks_df['index'] == RIGHT_EYE].iloc[0]
|
||||
|
||||
dx = right_eye['x'] - left_eye['x']
|
||||
dy = right_eye['y'] - left_eye['y']
|
||||
return np.degrees(np.arctan2(dy, dx))
|
||||
|
||||
@staticmethod
|
||||
def rotate_image(image: Union[torch.Tensor, np.ndarray, Image.Image], landmarks_df: pd.DataFrame) -> Tuple[Optional[np.ndarray], Optional[pd.DataFrame]]:
|
||||
"""Rotate image based on facial landmarks."""
|
||||
image_np = ImageProcessor._convert_to_numpy(image)
|
||||
|
||||
if image_np is None or landmarks_df is None:
|
||||
return None, None
|
||||
|
||||
angle = ImageProcessor.calculate_rotation_angle(landmarks_df)
|
||||
height, width = image_np.shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
|
||||
rotated_image = cv2.warpAffine(image_np, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Transform landmarks
|
||||
ones = np.ones(shape=(len(landmarks_df), 1))
|
||||
points = np.hstack([landmarks_df[['x', 'y']].values, ones])
|
||||
transformed_points = rotation_matrix.dot(points.T).T
|
||||
|
||||
updated_landmarks = landmarks_df.copy()
|
||||
updated_landmarks['x'] = transformed_points[:, 0]
|
||||
updated_landmarks['y'] = transformed_points[:, 1]
|
||||
|
||||
return rotated_image, updated_landmarks
|
||||
|
||||
@staticmethod
|
||||
def crop_face_to_square(image: np.ndarray, landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Tuple[
|
||||
Optional[np.ndarray], Optional[Tuple[int, int, int, int]]]:
|
||||
"""Crop face region to a square (1:1) based on landmarks."""
|
||||
bbox = ImageProcessor.calculate_face_bbox(landmarks_df, padding_percent)
|
||||
if bbox is None:
|
||||
return None, None
|
||||
|
||||
x, y, w, h = bbox
|
||||
|
||||
# Calculate the center of the bounding box
|
||||
center_x = x + w // 2
|
||||
center_y = y + h // 2
|
||||
|
||||
# Determine the size of the square crop
|
||||
crop_size = max(w, h)
|
||||
half_size = crop_size // 2
|
||||
|
||||
# Calculate the crop coordinates
|
||||
x1 = max(0, center_x - half_size)
|
||||
y1 = max(0, center_y - half_size)
|
||||
x2 = min(image.shape[1], center_x + half_size)
|
||||
y2 = min(image.shape[0], center_y + half_size)
|
||||
|
||||
# Adjust if the crop goes out of bounds
|
||||
if x2 - x1 < crop_size:
|
||||
x1 = max(0, x2 - crop_size)
|
||||
if y2 - y1 < crop_size:
|
||||
y1 = max(0, y2 - crop_size)
|
||||
|
||||
# Crop the image
|
||||
cropped_face = image[y1:y2, x1:x2]
|
||||
|
||||
# Return the cropped face and the crop bounding box
|
||||
return cropped_face, (x1, y1, x2 - x1, y2 - y1)
|
||||
@@ -0,0 +1,329 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
import pandas as pd
|
||||
from typing import Union, Optional, Tuple
|
||||
|
||||
from core.face_detector import FaceDetector
|
||||
|
||||
|
||||
class FaceFitAndRestore:
|
||||
"""ComfyUI node for fitting and restoring faces with a mode selection."""
|
||||
|
||||
def __init__(self):
|
||||
self.face_detector = FaceDetector()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mode": (["Fit", "Restore"], {
|
||||
"default": "Fit"
|
||||
}),
|
||||
"image": ("IMAGE",),
|
||||
"padding_percent": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05
|
||||
}),
|
||||
"bbox_size": (["512", "1024", "2048"], {
|
||||
"default": "1024"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"face_settings": ("FACE_SETTINGS", {
|
||||
"default": None
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "FACE_SETTINGS", "MASK")
|
||||
RETURN_NAMES = ("image", "face_settings", "mask")
|
||||
FUNCTION = "process_image"
|
||||
CATEGORY = "Face Processor"
|
||||
|
||||
def process_image(self, mode, image, padding_percent=0.0, bbox_size="1024", face_settings=None):
|
||||
if mode == "Fit":
|
||||
return self._fit(image, padding_percent, bbox_size)
|
||||
elif mode == "Restore":
|
||||
if face_settings is None:
|
||||
print("Face settings are required in Restore mode, returning original image")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
return self._restore(image, face_settings)
|
||||
else:
|
||||
print(f"Invalid mode: {mode}, returning original image")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
def _fit(self, image, padding_percent, bbox_size):
|
||||
"""Fit mode: Crop and process the face."""
|
||||
# Convert tensor to numpy
|
||||
image_np = self._convert_to_numpy(image)
|
||||
if image_np is None:
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Detect facial landmarks
|
||||
landmarks_df = self.face_detector.detect_landmarks(image_np)
|
||||
if landmarks_df is None:
|
||||
print("No face detected, returning original image")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Calculate rotation angle and rotate the image
|
||||
rotation_angle = self._calculate_rotation_angle(landmarks_df)
|
||||
rotated_image, updated_landmarks = self._rotate_image(image_np, landmarks_df)
|
||||
if rotated_image is None:
|
||||
print("Failed to rotate image, returning original")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Crop face region to a square (1:1)
|
||||
cropped_face, crop_bbox = self._crop_face_to_square(rotated_image, updated_landmarks, padding_percent)
|
||||
if cropped_face is None:
|
||||
print("Failed to crop face, returning original image")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Resize to target size
|
||||
target_size = int(bbox_size)
|
||||
final_image = self._resize_image(cropped_face, target_size)
|
||||
if final_image is None:
|
||||
print("Failed to resize image, returning original")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Convert back to float32 format expected by ComfyUI
|
||||
final_image = final_image.astype(np.float32) / 255.0
|
||||
final_image = torch.from_numpy(final_image).unsqueeze(0)
|
||||
|
||||
# Save face settings for restoration
|
||||
face_settings = {
|
||||
"original_image_shape": image_np.shape,
|
||||
"rotation_angle": rotation_angle,
|
||||
"crop_bbox": crop_bbox, # (x, y, w, h)
|
||||
"padding_percent": padding_percent,
|
||||
"bbox_size": target_size,
|
||||
}
|
||||
|
||||
# Create a mask for the face area (1:1)
|
||||
# Initialize mask with ones (all pixels are part of the face)
|
||||
mask = np.ones((target_size, target_size), dtype=np.float32)
|
||||
|
||||
# If there was a rotation, exclude areas that are outside the original image
|
||||
if rotation_angle != 0:
|
||||
# Create a blank image with the original size
|
||||
original_mask = np.ones(image_np.shape[:2], dtype=np.float32)
|
||||
|
||||
# Rotate the original mask to match the rotated image
|
||||
height, width = image_np.shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, rotation_angle, 1.0)
|
||||
rotated_mask = cv2.warpAffine(original_mask, rotation_matrix, (width, height), flags=cv2.INTER_LINEAR)
|
||||
|
||||
# Crop the rotated mask to the face region
|
||||
x1, y1, w, h = crop_bbox
|
||||
cropped_mask = rotated_mask[y1:y1 + h, x1:x1 + w]
|
||||
|
||||
# Resize the cropped mask to the target size
|
||||
resized_mask = cv2.resize(cropped_mask, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Update the final mask to exclude areas outside the original image
|
||||
mask = resized_mask
|
||||
|
||||
# Convert mask to torch tensor
|
||||
mask = torch.from_numpy(mask).unsqueeze(0)
|
||||
|
||||
return (final_image, face_settings, mask)
|
||||
|
||||
def _restore(self, image, face_settings):
|
||||
"""Restore mode: Restore the face to the original image."""
|
||||
# Convert tensor to numpy
|
||||
processed_face_np = self._convert_to_numpy(image)
|
||||
if processed_face_np is None:
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Extract face settings
|
||||
original_image_shape = face_settings.get("original_image_shape")
|
||||
rotation_angle = face_settings.get("rotation_angle")
|
||||
crop_bbox = face_settings.get("crop_bbox")
|
||||
padding_percent = face_settings.get("padding_percent")
|
||||
bbox_size = face_settings.get("bbox_size")
|
||||
|
||||
if not all([original_image_shape, crop_bbox]):
|
||||
print("Invalid face settings, returning processed face")
|
||||
return (image, {}, self._create_empty_mask(image))
|
||||
|
||||
# Create a blank image with the original size
|
||||
restored_image = np.zeros(original_image_shape, dtype=np.uint8)
|
||||
|
||||
# Resize the processed face back to the original crop size
|
||||
x1, y1, w, h = crop_bbox
|
||||
resized_face = cv2.resize(processed_face_np, (w, h), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Place the resized face back into the rotated image
|
||||
restored_image[y1:y1 + h, x1:x1 + w] = resized_face
|
||||
|
||||
# Rotate the image back to the original orientation (reverse the rotation)
|
||||
if rotation_angle != 0:
|
||||
height, width = original_image_shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0) # Reverse the rotation
|
||||
restored_image = cv2.warpAffine(restored_image, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Convert back to float32 format expected by ComfyUI
|
||||
restored_image = restored_image.astype(np.float32) / 255.0
|
||||
restored_image = torch.from_numpy(restored_image).unsqueeze(0)
|
||||
|
||||
# Create a mask for the face area (original image size)
|
||||
mask = self._create_mask(original_image_shape, crop_bbox, rotation_angle)
|
||||
|
||||
return (restored_image, {}, mask)
|
||||
|
||||
def _create_mask(self, image_shape, crop_bbox, rotation_angle):
|
||||
"""Create a mask for the face area."""
|
||||
mask = np.zeros(image_shape[:2], dtype=np.float32)
|
||||
x1, y1, w, h = crop_bbox
|
||||
|
||||
# Draw a white rectangle for the face area
|
||||
mask[y1:y1 + h, x1:x1 + w] = 1.0
|
||||
|
||||
# Rotate the mask back to the original orientation (reverse the rotation)
|
||||
if rotation_angle != 0:
|
||||
height, width = image_shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0)
|
||||
mask = cv2.warpAffine(mask, rotation_matrix, (width, height), flags=cv2.INTER_LINEAR)
|
||||
|
||||
# Convert to torch tensor
|
||||
mask = torch.from_numpy(mask).unsqueeze(0)
|
||||
return mask
|
||||
|
||||
def _create_empty_mask(self, image):
|
||||
"""Create an empty mask with the same shape as the input image."""
|
||||
if torch.is_tensor(image):
|
||||
image_shape = image.shape[1:3] # (H, W)
|
||||
else:
|
||||
image_shape = image.shape[:2] # (H, W)
|
||||
mask = torch.zeros((1, *image_shape), dtype=torch.float32)
|
||||
return mask
|
||||
|
||||
def _convert_to_numpy(self, image: torch.Tensor) -> Optional[np.ndarray]:
|
||||
"""Convert tensor to numpy array."""
|
||||
if torch.is_tensor(image):
|
||||
image = image.detach().cpu().numpy()
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
image = (image * 255).astype(np.uint8)
|
||||
return image
|
||||
return None
|
||||
|
||||
def _calculate_rotation_angle(self, landmarks_df: pd.DataFrame) -> float:
|
||||
"""Calculate rotation angle based on eyes position."""
|
||||
if landmarks_df is None or landmarks_df.empty:
|
||||
return 0.0
|
||||
|
||||
LEFT_EYE = 33 # Center of the left eye
|
||||
RIGHT_EYE = 263 # Center of the right eye
|
||||
|
||||
left_eye = landmarks_df[landmarks_df['index'] == LEFT_EYE].iloc[0]
|
||||
right_eye = landmarks_df[landmarks_df['index'] == RIGHT_EYE].iloc[0]
|
||||
|
||||
dx = right_eye['x'] - left_eye['x']
|
||||
dy = right_eye['y'] - left_eye['y']
|
||||
return np.degrees(np.arctan2(dy, dx))
|
||||
|
||||
def _rotate_image(self, image: np.ndarray, landmarks_df: pd.DataFrame) -> Tuple[Optional[np.ndarray], Optional[pd.DataFrame]]:
|
||||
"""Rotate image based on facial landmarks."""
|
||||
if image is None or landmarks_df is None:
|
||||
return None, None
|
||||
|
||||
angle = self._calculate_rotation_angle(landmarks_df)
|
||||
height, width = image.shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
|
||||
rotated_image = cv2.warpAffine(image, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Transform landmarks
|
||||
ones = np.ones(shape=(len(landmarks_df), 1))
|
||||
points = np.hstack([landmarks_df[['x', 'y']].values, ones])
|
||||
transformed_points = rotation_matrix.dot(points.T).T
|
||||
|
||||
updated_landmarks = landmarks_df.copy()
|
||||
updated_landmarks['x'] = transformed_points[:, 0]
|
||||
updated_landmarks['y'] = transformed_points[:, 1]
|
||||
|
||||
return rotated_image, updated_landmarks
|
||||
|
||||
def _crop_face_to_square(self, image: np.ndarray, landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Tuple[Optional[np.ndarray], Optional[Tuple[int, int, int, int]]]:
|
||||
"""Crop face region to a square (1:1) based on landmarks."""
|
||||
bbox = self._calculate_face_bbox(landmarks_df, padding_percent)
|
||||
if bbox is None:
|
||||
return None, None
|
||||
|
||||
x, y, w, h = bbox
|
||||
|
||||
# Calculate the center of the bounding box
|
||||
center_x = x + w // 2
|
||||
center_y = y + h // 2
|
||||
|
||||
# Determine the size of the square crop
|
||||
crop_size = max(w, h)
|
||||
half_size = crop_size // 2
|
||||
|
||||
# Calculate the crop coordinates
|
||||
x1 = max(0, center_x - half_size)
|
||||
y1 = max(0, center_y - half_size)
|
||||
x2 = min(image.shape[1], center_x + half_size)
|
||||
y2 = min(image.shape[0], center_y + half_size)
|
||||
|
||||
# Adjust if the crop goes out of bounds
|
||||
if x2 - x1 < crop_size:
|
||||
x1 = max(0, x2 - crop_size)
|
||||
if y2 - y1 < crop_size:
|
||||
y1 = max(0, y2 - crop_size)
|
||||
|
||||
# Crop the image
|
||||
cropped_face = image[y1:y2, x1:x2]
|
||||
|
||||
# Return the cropped face and the crop bounding box
|
||||
return cropped_face, (x1, y1, x2 - x1, y2 - y1)
|
||||
|
||||
def _calculate_face_bbox(self, landmarks_df: pd.DataFrame, padding_percent: float = 0.0) -> Optional[Tuple[int, int, int, int]]:
|
||||
"""Calculate bounding box for face based on landmarks."""
|
||||
if landmarks_df is None or landmarks_df.empty:
|
||||
return None
|
||||
|
||||
min_x = landmarks_df['x'].min()
|
||||
max_x = landmarks_df['x'].max()
|
||||
min_y = landmarks_df['y'].min()
|
||||
max_y = landmarks_df['y'].max()
|
||||
|
||||
width = max_x - min_x
|
||||
height = max_y - min_y
|
||||
|
||||
pad_x = width * padding_percent
|
||||
pad_y = height * padding_percent
|
||||
|
||||
x1 = max(0, min_x - pad_x)
|
||||
y1 = max(0, min_y - pad_y)
|
||||
x2 = max_x + pad_x
|
||||
y2 = max_y + pad_y
|
||||
|
||||
return (int(x1), int(y1), int(x2 - x1), int(y2 - y1))
|
||||
|
||||
def _resize_image(self, image: np.ndarray, target_size: int) -> Optional[np.ndarray]:
|
||||
"""Resize image to target size while maintaining aspect ratio."""
|
||||
if image is None:
|
||||
return None
|
||||
|
||||
h, w = image.shape[:2]
|
||||
scale = min(target_size / h, target_size / w)
|
||||
new_h = int(h * scale)
|
||||
new_w = int(w * scale)
|
||||
|
||||
resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
square = np.zeros((target_size, target_size, image.shape[2]), dtype=resized.dtype)
|
||||
y_offset = (target_size - new_h) // 2
|
||||
x_offset = (target_size - new_w) // 2
|
||||
|
||||
square[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized
|
||||
|
||||
return square
|
||||
@@ -0,0 +1,94 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
import pandas as pd
|
||||
from typing import Union, Optional, Tuple
|
||||
|
||||
from core.face_detector import FaceDetector
|
||||
from core.image_processor import ImageProcessor
|
||||
|
||||
|
||||
class FitAndCrop:
|
||||
"""ComfyUI node for detecting, cropping, and fitting faces to a square bbox."""
|
||||
|
||||
def __init__(self):
|
||||
self.face_detector = FaceDetector()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"padding_percent": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05
|
||||
}),
|
||||
"bbox_size": (["512", "1024", "2048"], {
|
||||
"default": "1024"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "FACE_SETTINGS")
|
||||
RETURN_NAMES = ("image", "face_settings")
|
||||
FUNCTION = "process_image"
|
||||
CATEGORY = "Face Processor"
|
||||
|
||||
def process_image(self, image, padding_percent=0.0, bbox_size="1024"):
|
||||
# Convert tensor to numpy
|
||||
image_np = self._convert_to_numpy(image)
|
||||
if image_np is None:
|
||||
return (image, {})
|
||||
|
||||
# Detect facial landmarks
|
||||
landmarks_df = self.face_detector.detect_landmarks(image_np)
|
||||
if landmarks_df is None:
|
||||
print("No face detected, returning original image")
|
||||
return (image, {})
|
||||
|
||||
# Calculate rotation angle and rotate the image
|
||||
rotation_angle = ImageProcessor.calculate_rotation_angle(landmarks_df)
|
||||
rotated_image, updated_landmarks = ImageProcessor.rotate_image(image_np, landmarks_df)
|
||||
if rotated_image is None:
|
||||
print("Failed to rotate image, returning original")
|
||||
return (image, {})
|
||||
|
||||
# Crop face region to a square (1:1)
|
||||
cropped_face, crop_bbox = ImageProcessor.crop_face_to_square(rotated_image, updated_landmarks, padding_percent)
|
||||
if cropped_face is None:
|
||||
print("Failed to crop face, returning original image")
|
||||
return (image, {})
|
||||
|
||||
# Resize to target size
|
||||
target_size = int(bbox_size)
|
||||
final_image = ImageProcessor.resize_image(cropped_face, target_size)
|
||||
if final_image is None:
|
||||
print("Failed to resize image, returning original")
|
||||
return (image, {})
|
||||
|
||||
# Convert back to float32 format expected by ComfyUI
|
||||
final_image = final_image.astype(np.float32) / 255.0
|
||||
final_image = torch.from_numpy(final_image).unsqueeze(0)
|
||||
|
||||
# Save face settings for restoration
|
||||
face_settings = {
|
||||
"original_image_shape": image_np.shape,
|
||||
"rotation_angle": rotation_angle,
|
||||
"crop_bbox": crop_bbox, # (x, y, w, h)
|
||||
"padding_percent": padding_percent,
|
||||
"bbox_size": target_size,
|
||||
}
|
||||
|
||||
return (final_image, face_settings)
|
||||
|
||||
def _convert_to_numpy(self, image: torch.Tensor) -> Optional[np.ndarray]:
|
||||
"""Convert tensor to numpy array."""
|
||||
if torch.is_tensor(image):
|
||||
image = image.detach().cpu().numpy()
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
image = (image * 255).astype(np.uint8)
|
||||
return image
|
||||
return None
|
||||
@@ -0,0 +1,70 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
from typing import Union, Optional
|
||||
|
||||
class RestoreFace:
|
||||
"""ComfyUI node to restore a processed face back to its original position in the image."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"processed_face": ("IMAGE",),
|
||||
"face_settings": ("FACE_SETTINGS",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "restore_face"
|
||||
CATEGORY = "Face Processor"
|
||||
|
||||
def restore_face(self, processed_face, face_settings):
|
||||
# Convert processed face to numpy
|
||||
processed_face_np = self._convert_to_numpy(processed_face)
|
||||
if processed_face_np is None:
|
||||
return (processed_face,)
|
||||
|
||||
# Extract face settings
|
||||
original_image_shape = face_settings.get("original_image_shape")
|
||||
rotation_angle = face_settings.get("rotation_angle")
|
||||
crop_bbox = face_settings.get("crop_bbox")
|
||||
padding_percent = face_settings.get("padding_percent")
|
||||
bbox_size = face_settings.get("bbox_size")
|
||||
|
||||
if not all([original_image_shape, crop_bbox]):
|
||||
print("Invalid face settings, returning processed face")
|
||||
return (processed_face,)
|
||||
|
||||
# Resize the processed face back to the original crop size
|
||||
x1, y1, w, h = crop_bbox
|
||||
resized_face = cv2.resize(processed_face_np, (w, h), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Create a blank image with the original size
|
||||
restored_image = np.zeros(original_image_shape, dtype=np.uint8)
|
||||
|
||||
# Place the resized face back into the rotated image
|
||||
restored_image[y1:y1 + h, x1:x1 + w] = resized_face
|
||||
|
||||
# Rotate the image back to the original orientation (reverse the rotation)
|
||||
if rotation_angle != 0:
|
||||
height, width = original_image_shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0) # Reverse the rotation
|
||||
restored_image = cv2.warpAffine(restored_image, rotation_matrix, (width, height), flags=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Convert back to float32 format expected by ComfyUI
|
||||
restored_image = restored_image.astype(np.float32) / 255.0
|
||||
restored_image = torch.from_numpy(restored_image).unsqueeze(0)
|
||||
|
||||
return (restored_image,)
|
||||
|
||||
def _convert_to_numpy(self, image: torch.Tensor) -> Optional[np.ndarray]:
|
||||
"""Convert tensor to numpy array."""
|
||||
if torch.is_tensor(image):
|
||||
image = image.detach().cpu().numpy()
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
image = (image * 255).astype(np.uint8)
|
||||
return image
|
||||
return None
|
||||
Reference in New Issue
Block a user