- Face Fit or Restore fixes
WIP - base landmarks and obj file
This commit is contained in:
Rafał Tłuchowski
+1
-1
@@ -18,4 +18,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FaceFitAndRestore": "Face Fit or Restore"
|
||||
}
|
||||
|
||||
__version__ = "1.0.0"
|
||||
__version__ = "1.0.1"
|
||||
+2
-2
@@ -57,10 +57,10 @@ def generate_context_file(
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
output_file = "D:/code/code_base/face_processor_02.txt"
|
||||
output_file = "D:/code/code_base/face_processor_03.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"},
|
||||
ignore_files={".gitignore", "contex_generator.py", output_file, "package-lock.json", "TODO.txt", "mediapipe_landmarks_468.obj"},
|
||||
)
|
||||
@@ -0,0 +1,288 @@
|
||||
import numpy as np
|
||||
import os
|
||||
|
||||
|
||||
class MediapipeBaseLandmarks:
|
||||
"""
|
||||
Class for handling base landmark positions and face topology loaded from OBJ file.
|
||||
Coordinates are stored in normalized 0-1 range.
|
||||
"""
|
||||
_base_landmarks = None
|
||||
_faces = None
|
||||
_boundary_faces = None
|
||||
_boundary_landmarks = None
|
||||
|
||||
# Face oval landmarks in clockwise order
|
||||
OVAL_LANDMARKS = [10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288,
|
||||
397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136,
|
||||
172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109]
|
||||
|
||||
@classmethod
|
||||
def _transform_landmarks(cls, landmarks, x_scale=1.0, y_translation=0.0):
|
||||
"""
|
||||
Apply transformations to landmarks:
|
||||
- x_scale: horizontal scaling (0.5 to 1.0)
|
||||
- y_translation: vertical translation (-0.5 to 0.5)
|
||||
|
||||
Args:
|
||||
landmarks: numpy array of landmarks
|
||||
x_scale: horizontal scaling factor (default: 1.0)
|
||||
y_translation: vertical translation value (default: 0.0)
|
||||
|
||||
Returns:
|
||||
numpy array: Transformed landmarks
|
||||
"""
|
||||
transformed = landmarks.copy()
|
||||
|
||||
# Calculate center point for scaling
|
||||
center_x = (landmarks[:, 0].min() + landmarks[:, 0].max()) / 2
|
||||
|
||||
# Apply horizontal scaling relative to center
|
||||
transformed[:, 0] = center_x + (transformed[:, 0] - center_x) * x_scale
|
||||
|
||||
# Apply vertical translation
|
||||
transformed[:, 1] += y_translation
|
||||
|
||||
return transformed
|
||||
|
||||
@classmethod
|
||||
def _load_obj_data(cls):
|
||||
"""
|
||||
Load UV coordinates, vertices and face indices from OBJ file.
|
||||
Returns array of normalized [u,v] coordinates mapped to correct vertex indices
|
||||
and array of face indices.
|
||||
"""
|
||||
if cls._base_landmarks is not None and cls._faces is not None:
|
||||
return cls._base_landmarks, cls._faces
|
||||
|
||||
current_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
project_root = os.path.dirname(current_dir)
|
||||
obj_path = os.path.join(project_root, 'models', 'uv_maps', 'mediapipe_landmarks_468.obj')
|
||||
|
||||
vertices = []
|
||||
uvs = []
|
||||
faces = []
|
||||
vertex_to_uv = {}
|
||||
|
||||
print(f"Loading 3D face model from: {obj_path}")
|
||||
|
||||
with open(obj_path, 'r') as f:
|
||||
for line in f:
|
||||
if line.startswith('v '):
|
||||
parts = line.strip().split()
|
||||
vertices.append([float(p) for p in parts[1:4]])
|
||||
elif line.startswith('vt '):
|
||||
parts = line.strip().split()
|
||||
u = float(parts[1])
|
||||
v = 1.0 - float(parts[2])
|
||||
uvs.append([u, v])
|
||||
elif line.startswith('f '):
|
||||
parts = line.strip().split()[1:]
|
||||
face_verts = []
|
||||
for part in parts:
|
||||
indices = part.split('/')
|
||||
if len(indices) >= 2:
|
||||
vertex_idx = int(indices[0]) - 1
|
||||
uv_idx = int(indices[1]) - 1
|
||||
vertex_to_uv[vertex_idx] = uv_idx
|
||||
face_verts.append(vertex_idx)
|
||||
if len(face_verts) == 3:
|
||||
faces.append(face_verts)
|
||||
|
||||
print(f"Loaded {len(vertices)} vertices, {len(uvs)} UVs, {len(faces)} faces")
|
||||
|
||||
uvs = np.array(uvs, dtype=np.float32)
|
||||
faces = np.array(faces, dtype=np.int32)
|
||||
landmarks = np.zeros((468, 2), dtype=np.float32)
|
||||
|
||||
for vertex_idx, uv_idx in vertex_to_uv.items():
|
||||
if vertex_idx < len(landmarks):
|
||||
landmarks[vertex_idx] = uvs[uv_idx]
|
||||
|
||||
cls._base_landmarks = landmarks
|
||||
cls._faces = faces
|
||||
return cls._base_landmarks, cls._faces
|
||||
|
||||
@staticmethod
|
||||
def interpolate_edge_points(start_pos, end_pos, n_segments):
|
||||
"""
|
||||
Interpolate points along an edge
|
||||
"""
|
||||
points = []
|
||||
for i in range(n_segments + 1):
|
||||
t = i / n_segments
|
||||
x = round(start_pos[0] + t * (end_pos[0] - start_pos[0]))
|
||||
y = round(start_pos[1] + t * (end_pos[1] - start_pos[1]))
|
||||
points.append((x, y))
|
||||
return points
|
||||
|
||||
@classmethod
|
||||
def create_boundary_triangles(cls, size):
|
||||
"""
|
||||
Create additional triangles between face oval and image boundaries
|
||||
"""
|
||||
if isinstance(size, int):
|
||||
width = height = size
|
||||
else:
|
||||
width, height = size
|
||||
|
||||
print(f"Creating boundary triangles for image size: {width}x{height}")
|
||||
|
||||
boundary_points = []
|
||||
boundary_triangles = []
|
||||
|
||||
edges = {
|
||||
'top': {
|
||||
'landmarks': [54, 103, 67, 109, 10, 338, 297, 332, 284],
|
||||
'start': (0, 0),
|
||||
'middle': (width / 2, 0),
|
||||
'end': (width, 0),
|
||||
'segments': 8
|
||||
},
|
||||
'right': {
|
||||
'landmarks': [284, 251, 389, 356, 454, 323, 361, 288, 397, 365],
|
||||
'start': (width, 0),
|
||||
'middle': (width, height / 2),
|
||||
'end': (width, height),
|
||||
'segments': 9
|
||||
},
|
||||
'bottom': {
|
||||
'landmarks': [365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136],
|
||||
'start': (width, height),
|
||||
'middle': (width / 2, height),
|
||||
'end': (0, height),
|
||||
'segments': 10
|
||||
},
|
||||
'left': {
|
||||
'landmarks': [136, 172, 58, 132, 93, 234, 127, 162, 21, 54],
|
||||
'start': (0, height),
|
||||
'middle': (0, height / 2),
|
||||
'end': (0, 0),
|
||||
'segments': 9
|
||||
}
|
||||
}
|
||||
|
||||
current_point_idx = len(cls._base_landmarks)
|
||||
|
||||
for edge_name, edge_info in edges.items():
|
||||
print(f"Processing {edge_name} edge")
|
||||
|
||||
landmarks = edge_info['landmarks']
|
||||
n_segments = edge_info['segments']
|
||||
print(f"Creating {n_segments} segments for {len(landmarks)} landmarks")
|
||||
|
||||
# Generate edge points
|
||||
edge_points = cls.interpolate_edge_points(
|
||||
edge_info['start'],
|
||||
edge_info['end'],
|
||||
n_segments
|
||||
)
|
||||
|
||||
# Create quads and convert to triangles
|
||||
for i in range(n_segments):
|
||||
# Define quad vertices
|
||||
quad = [
|
||||
landmarks[i], # Current landmark
|
||||
landmarks[i + 1], # Next landmark
|
||||
current_point_idx + i + 1, # Next edge point
|
||||
current_point_idx + i # Current edge point
|
||||
]
|
||||
|
||||
# Convert quad to two triangles
|
||||
boundary_triangles.extend([
|
||||
[quad[0], quad[1], quad[2]], # First triangle
|
||||
[quad[0], quad[2], quad[3]] # Second triangle
|
||||
])
|
||||
|
||||
# Add edge points to boundary points list
|
||||
boundary_points.extend(edge_points)
|
||||
current_point_idx += len(edge_points)
|
||||
|
||||
print(f"Created {len(boundary_triangles)} triangles from {len(boundary_points)} boundary points")
|
||||
|
||||
cls._boundary_landmarks = np.array(boundary_points)
|
||||
cls._boundary_faces = np.array(boundary_triangles)
|
||||
|
||||
return cls._boundary_faces
|
||||
|
||||
@classmethod
|
||||
def get_base_landmarks(cls, size=None, x_scale=1.0, y_translation=0.0):
|
||||
"""
|
||||
Get base landmark positions with optional transformations.
|
||||
|
||||
Args:
|
||||
size: Target size (int or tuple)
|
||||
x_scale: Horizontal scaling factor (0.5 to 1.0)
|
||||
y_translation: Vertical translation (-0.5 to 0.5)
|
||||
|
||||
Returns:
|
||||
numpy array: Base landmarks with applied transformations
|
||||
"""
|
||||
landmarks, _ = cls._load_obj_data()
|
||||
|
||||
# Apply transformations before scaling to size
|
||||
landmarks = cls._transform_landmarks(landmarks, x_scale, y_translation)
|
||||
|
||||
if size is not None:
|
||||
if isinstance(size, tuple):
|
||||
landmarks = landmarks * np.array(size)
|
||||
else:
|
||||
landmarks = landmarks * size
|
||||
|
||||
return landmarks
|
||||
|
||||
@classmethod
|
||||
def get_face_triangles(cls, size=None, x_scale=1.0, y_translation=0.0):
|
||||
"""
|
||||
Get face triangulation including boundary triangles with optional transformations
|
||||
|
||||
Args:
|
||||
size: Target size (int or tuple)
|
||||
x_scale: Horizontal scaling factor (0.5 to 1.0)
|
||||
y_translation: Vertical translation (-0.5 to 0.5)
|
||||
|
||||
Returns:
|
||||
tuple: (triangles array, landmarks array)
|
||||
"""
|
||||
landmarks, faces = cls._load_obj_data()
|
||||
|
||||
# Apply transformations
|
||||
landmarks = cls._transform_landmarks(landmarks, x_scale, y_translation)
|
||||
|
||||
if size is not None:
|
||||
if isinstance(size, tuple):
|
||||
scaled_landmarks = landmarks * np.array(size)
|
||||
else:
|
||||
scaled_landmarks = landmarks * size
|
||||
|
||||
if cls._boundary_faces is None:
|
||||
cls.create_boundary_triangles(size)
|
||||
|
||||
if cls._boundary_faces is not None:
|
||||
# Transform boundary landmarks if they exist
|
||||
if cls._boundary_landmarks is not None:
|
||||
boundary_landmarks = cls._boundary_landmarks
|
||||
# Note: We don't apply x_scale to boundary landmarks as they should
|
||||
# stay at the image edges, but we do apply y_translation
|
||||
if y_translation != 0:
|
||||
boundary_landmarks = boundary_landmarks.copy()
|
||||
boundary_landmarks[:, 1] += y_translation * (size if isinstance(size, int) else size[1])
|
||||
|
||||
all_points = np.vstack([scaled_landmarks, boundary_landmarks])
|
||||
all_triangles = np.vstack([faces, cls._boundary_faces])
|
||||
return all_triangles, all_points
|
||||
|
||||
return faces, landmarks
|
||||
|
||||
@staticmethod
|
||||
def validate_size(size):
|
||||
"""
|
||||
Validate size parameter
|
||||
"""
|
||||
if isinstance(size, (int, tuple)):
|
||||
if isinstance(size, int) and size <= 0:
|
||||
raise ValueError("Size must be a positive integer")
|
||||
if isinstance(size, tuple) and (len(size) != 2 or size[0] <= 0 or size[1] <= 0):
|
||||
raise ValueError("Size tuple must contain two positive integers")
|
||||
return size
|
||||
raise ValueError("Size must be an integer or tuple of two integers")
|
||||
@@ -167,4 +167,4 @@ class ImageProcessor:
|
||||
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)
|
||||
return cropped_face, (x1, y1, x2 - x1, y2 - y1)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,94 +0,0 @@
|
||||
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
|
||||
@@ -1,70 +0,0 @@
|
||||
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