fp_pipe normalization across different workflows
This commit is contained in:
Rafał Tłuchowski
+203
-153
@@ -8,6 +8,8 @@ from ..core.image_processor import ImageProcessor
|
||||
|
||||
|
||||
class FaceFitAndRestore:
|
||||
"""ComfyUI node for processing face images in Fit or Restore mode."""
|
||||
|
||||
def __init__(self):
|
||||
self.face_detector = FaceDetector()
|
||||
self.image_processor = ImageProcessor()
|
||||
@@ -50,245 +52,293 @@ class FaceFitAndRestore:
|
||||
|
||||
def process_image(self, mode, workflow, padding_percent=0.0, bbox_size="1024",
|
||||
image=None, fp_pipe=None, image_sequence=None):
|
||||
"""Process images in either Fit or Restore mode with unified fp_pipe structure."""
|
||||
|
||||
# Validate inputs
|
||||
if image is None and (workflow == "single_image" or mode == "Restore"):
|
||||
print("Error: Image is required for single_image workflow and Restore mode")
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
# Check if we should use workflow from fp_pipe in Restore mode
|
||||
if mode == "Restore":
|
||||
if fp_pipe is None:
|
||||
print("Error: fp_pipe is required in Restore mode")
|
||||
return None, {}, None, int(bbox_size)
|
||||
if workflow == "image_sequence" and mode == "Fit" and (not image_sequence or "frames" not in image_sequence):
|
||||
print("Error: Valid image sequence data is required for Fit mode")
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
if "workflow" in fp_pipe:
|
||||
workflow = fp_pipe["workflow"]
|
||||
print(f"Using workflow from pipe: {workflow}")
|
||||
if mode == "Restore" and (not fp_pipe or "frames" not in fp_pipe):
|
||||
print("Error: Valid fp_pipe with frames is required for Restore mode")
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
# Additional check for sequence restore
|
||||
if workflow == "image_sequence" and "frames" not in fp_pipe:
|
||||
print("Error: fp_pipe does not contain frames data required for image sequence restore")
|
||||
return None, {}, None, int(bbox_size)
|
||||
# Initialize fp_pipe if None
|
||||
if fp_pipe is None:
|
||||
fp_pipe = {
|
||||
"workflow": workflow,
|
||||
"current_frame": 0,
|
||||
"padding_percent": padding_percent,
|
||||
"target_lm": {}, # Will be populated by face_wrapper
|
||||
"frames": {}
|
||||
}
|
||||
else:
|
||||
# Update common parameters
|
||||
fp_pipe["workflow"] = workflow
|
||||
fp_pipe["padding_percent"] = padding_percent
|
||||
|
||||
if workflow == "single_image":
|
||||
if image is None:
|
||||
print("Error: Image is required for single_image workflow")
|
||||
return None, {}, None, int(bbox_size)
|
||||
try:
|
||||
if workflow == "single_image":
|
||||
return self._process_frame(
|
||||
mode=mode,
|
||||
image=image,
|
||||
frame_number=0,
|
||||
padding_percent=padding_percent,
|
||||
bbox_size=bbox_size,
|
||||
fp_pipe=fp_pipe
|
||||
)
|
||||
else: # image_sequence
|
||||
# W trybie Restore używamy current_frame z fp_pipe
|
||||
current_frame = (fp_pipe.get("current_frame", 0) if mode == "Restore"
|
||||
else image_sequence.get("current_frame", 0))
|
||||
fp_pipe["current_frame"] = current_frame
|
||||
|
||||
if mode == "Fit":
|
||||
result = self._fit(image, padding_percent, bbox_size)
|
||||
else: # Restore
|
||||
result = self._restore(image, fp_pipe)
|
||||
# W trybie Restore używamy liczby klatek z fp_pipe
|
||||
if mode == "Restore":
|
||||
total_frames = len(fp_pipe["frames"])
|
||||
frames_to_process = {
|
||||
idx: None for idx in range(total_frames)
|
||||
}
|
||||
else: # Fit mode
|
||||
total_frames = len(image_sequence["frames"])
|
||||
frames_to_process = image_sequence["frames"]
|
||||
|
||||
# Add workflow info to settings
|
||||
result_settings = result[1]
|
||||
result_settings["workflow"] = "single_image"
|
||||
return result
|
||||
|
||||
elif workflow == "image_sequence":
|
||||
# In Fit mode, we need image_sequence input
|
||||
if mode == "Fit":
|
||||
if image_sequence is None or not image_sequence.get("frames"):
|
||||
print("Error: Valid image_sequence is required for image_sequence workflow in Fit mode")
|
||||
return None, {}, None, int(bbox_size)
|
||||
sequence_settings = {
|
||||
"workflow": "image_sequence",
|
||||
"frames": {},
|
||||
"current_frame": image_sequence.get("current_frame", 0)
|
||||
}
|
||||
|
||||
selected_result = None
|
||||
total_frames = len(image_sequence["frames"])
|
||||
print(f"Processing sequence of {total_frames} frames...")
|
||||
print(f"Current frame index: {sequence_settings['current_frame']}")
|
||||
for frame_idx, frame_path in image_sequence["frames"].items():
|
||||
current_frame_result = None
|
||||
|
||||
# Process each frame
|
||||
for frame_idx in range(total_frames):
|
||||
print(f"Processing frame {frame_idx + 1}/{total_frames}")
|
||||
try:
|
||||
frame_tensor = self._load_image_from_path(frame_path)
|
||||
|
||||
if frame_tensor is None:
|
||||
print(f"Warning: Could not load frame {frame_idx} from {frame_path}")
|
||||
# Przygotowanie obrazu wejściowego
|
||||
if mode == "Fit":
|
||||
frame_path = frames_to_process[frame_idx]
|
||||
frame_image = self._load_image_from_path(frame_path)
|
||||
if frame_image is None:
|
||||
print(f"Warning: Could not load frame {frame_idx}")
|
||||
continue
|
||||
else: # Restore
|
||||
frame_image = image # Ten sam obraz dla każdej klatki w trybie Restore
|
||||
frame_path = None
|
||||
|
||||
result = self._fit(frame_tensor, padding_percent, bbox_size)
|
||||
# Process frame
|
||||
result = self._process_frame(
|
||||
mode=mode,
|
||||
image=frame_image,
|
||||
frame_number=frame_idx,
|
||||
padding_percent=padding_percent,
|
||||
bbox_size=bbox_size,
|
||||
fp_pipe=fp_pipe,
|
||||
original_path=frame_path if mode == "Fit" else None
|
||||
)
|
||||
|
||||
if result: # Check if result was created successfully
|
||||
result[1]["original_image_path"] = frame_path
|
||||
sequence_settings["frames"][f"frame_{frame_idx}"] = result[1]
|
||||
if frame_idx == sequence_settings["current_frame"]:
|
||||
selected_result = result
|
||||
# Store results for current frame
|
||||
if frame_idx == current_frame:
|
||||
current_frame_result = result
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing frame {frame_idx}: {str(e)}")
|
||||
continue
|
||||
# After processing all frames, return results for current frame
|
||||
if current_frame_result is None:
|
||||
print(f"Warning: Frame {current_frame} not found in sequence")
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
return current_frame_result
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in process_image: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
def _process_frame(self, mode, image, frame_number, padding_percent, bbox_size, fp_pipe, original_path=None):
|
||||
"""Process a single frame in either Fit or Restore mode."""
|
||||
try:
|
||||
frame_key = f"frame_{frame_number}"
|
||||
|
||||
if mode == "Fit":
|
||||
# Process frame in Fit mode
|
||||
result_image, frame_settings, result_mask = self._fit(
|
||||
image=image,
|
||||
padding_percent=padding_percent,
|
||||
bbox_size=bbox_size
|
||||
)
|
||||
|
||||
if frame_settings:
|
||||
# Create standardized frame data structure
|
||||
frame_data = {
|
||||
"bbox_size": int(bbox_size),
|
||||
"crop_bbox": frame_settings["crop_bbox"],
|
||||
"original_image_shape": frame_settings["original_image_shape"],
|
||||
"rotation_angle": frame_settings["rotation_angle"],
|
||||
"detected_lm": {} # Will be populated by face_wrapper
|
||||
}
|
||||
|
||||
# Add original path for sequence processing
|
||||
if original_path:
|
||||
frame_data["original_image_path"] = original_path
|
||||
|
||||
# Store frame data
|
||||
fp_pipe["frames"][frame_key] = frame_data
|
||||
|
||||
else: # Restore mode
|
||||
if image is None:
|
||||
print("Error: Input image is required for Restore mode")
|
||||
return None, {}, None, int(bbox_size)
|
||||
|
||||
current_frame = fp_pipe.get("current_frame", 0)
|
||||
frame_settings = fp_pipe.get("frames", {}).get(f"frame_{current_frame}")
|
||||
|
||||
frame_settings = fp_pipe["frames"].get(frame_key)
|
||||
if frame_settings is None:
|
||||
print(f"Error: No processor settings for frame {current_frame}")
|
||||
return None, {}, None, int(bbox_size)
|
||||
print(f"Processing restore for frame {current_frame} with settings: {frame_settings}")
|
||||
result = self._restore(image, frame_settings)
|
||||
|
||||
if result is None:
|
||||
print(f"Error: Failed to restore frame {current_frame}")
|
||||
print(f"Error: No settings found for frame {frame_number}")
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
return result[0], fp_pipe, result[2], int(bbox_size)
|
||||
|
||||
if selected_result is None:
|
||||
print(f"Warning: Could not process current frame, using default empty result")
|
||||
return None, sequence_settings, None, int(bbox_size)
|
||||
result = self._restore(image, frame_settings)
|
||||
if result is None:
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
return selected_result[0], sequence_settings, selected_result[2], int(bbox_size)
|
||||
result_image, result_mask = result[0], result[2]
|
||||
|
||||
return result_image, fp_pipe, result_mask, int(bbox_size)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing frame {frame_number}: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return None, fp_pipe, None, int(bbox_size)
|
||||
|
||||
def _fit(self, image, padding_percent, bbox_size):
|
||||
"""Fit mode: Crop and process the face."""
|
||||
# Convert image to numpy array
|
||||
image_np = self.image_processor.convert_to_numpy(image)
|
||||
if image_np is None:
|
||||
return image, {}, self._create_empty_mask(image), int(bbox_size)
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
# Detect facial landmarks using mediapipe
|
||||
# Detect facial landmarks
|
||||
landmarks_df = self.face_detector.detect_landmarks_mp(image_np)
|
||||
if landmarks_df is None:
|
||||
print("No face detected, returning original image")
|
||||
return image, {}, self._create_empty_mask(image), int(bbox_size)
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
# Calculate rotation angle and rotate the image
|
||||
# Calculate rotation and rotate image
|
||||
rotation_angle = self.image_processor.calculate_rotation_angle(landmarks_df)
|
||||
rotated_image, updated_landmarks = self.image_processor.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), int(bbox_size)
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
# Crop face region to a square (1:1)
|
||||
cropped_face, crop_bbox = self.image_processor.crop_face_to_square(rotated_image, updated_landmarks,
|
||||
padding_percent)
|
||||
# Crop face region
|
||||
cropped_face, crop_bbox = self.image_processor.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), int(bbox_size)
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
# Resize to target size
|
||||
target_size = int(bbox_size)
|
||||
final_image = self.image_processor.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), int(bbox_size)
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
# Convert to tensor format
|
||||
# Convert to tensor
|
||||
final_image = self.image_processor.numpy_to_tensor(final_image)
|
||||
|
||||
# Save processor settings for restoration
|
||||
fp_pipe = {
|
||||
# Create frame settings
|
||||
frame_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,
|
||||
"original_image_path": None
|
||||
"crop_bbox": crop_bbox,
|
||||
"bbox_size": target_size
|
||||
}
|
||||
|
||||
# Create a mask for the face area
|
||||
mask = np.ones((target_size, target_size), dtype=np.float32)
|
||||
# Create mask
|
||||
mask = self._create_mask(
|
||||
image_shape=frame_settings["original_image_shape"],
|
||||
crop_bbox=frame_settings["crop_bbox"],
|
||||
rotation_angle=frame_settings["rotation_angle"]
|
||||
)
|
||||
|
||||
if rotation_angle != 0:
|
||||
original_mask = np.ones(image_np.shape[:2], dtype=np.float32)
|
||||
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)
|
||||
|
||||
x1, y1, w, h = crop_bbox
|
||||
cropped_mask = rotated_mask[y1:y1 + h, x1:x1 + w]
|
||||
resized_mask = cv2.resize(cropped_mask, (target_size, target_size), interpolation=cv2.INTER_LINEAR)
|
||||
mask = resized_mask
|
||||
|
||||
mask = self.image_processor.convert_mask_to_tensor(mask)
|
||||
return final_image, fp_pipe, mask, int(bbox_size)
|
||||
|
||||
def _restore(self, image, fp_pipe):
|
||||
"""Restore mode: Restore the face to the original image."""
|
||||
|
||||
if fp_pipe is None:
|
||||
print("Error in restore: fp_pipe is None")
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
processed_face_np = self.image_processor.convert_to_numpy(image)
|
||||
if processed_face_np is None:
|
||||
print("Error in restore: Failed to convert image to numpy array")
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
|
||||
original_image_shape = fp_pipe.get("original_image_shape")
|
||||
rotation_angle = fp_pipe.get("rotation_angle")
|
||||
crop_bbox = fp_pipe.get("crop_bbox")
|
||||
padding_percent = fp_pipe.get("padding_percent")
|
||||
bbox_size = fp_pipe.get("bbox_size")
|
||||
|
||||
# Debug info
|
||||
print(f"Restore settings: shape={original_image_shape}, rotation={rotation_angle}, bbox={crop_bbox}")
|
||||
|
||||
if not all([original_image_shape, crop_bbox]):
|
||||
missing = []
|
||||
if not original_image_shape:
|
||||
missing.append("original_image_shape")
|
||||
if not crop_bbox:
|
||||
missing.append("crop_bbox")
|
||||
print(f"Error in restore: Missing required settings: {', '.join(missing)}")
|
||||
return image, {}, self._create_empty_mask(image)
|
||||
return final_image, frame_settings, mask
|
||||
|
||||
def _restore(self, image, frame_settings):
|
||||
"""Restore mode: Restore the face back to original position."""
|
||||
try:
|
||||
restored_image = np.zeros(original_image_shape, dtype=np.uint8)
|
||||
# Validate frame settings
|
||||
required_keys = ["original_image_shape", "rotation_angle", "crop_bbox"]
|
||||
if not all(key in frame_settings for key in required_keys):
|
||||
print("Error: Missing required frame settings")
|
||||
return None
|
||||
|
||||
x1, y1, w, h = crop_bbox
|
||||
# Convert input image to numpy
|
||||
processed_face_np = self.image_processor.convert_to_numpy(image)
|
||||
if processed_face_np is None:
|
||||
return None
|
||||
|
||||
# Get settings
|
||||
original_shape = frame_settings["original_image_shape"]
|
||||
rotation_angle = frame_settings["rotation_angle"]
|
||||
x1, y1, w, h = frame_settings["crop_bbox"]
|
||||
|
||||
# Create output image
|
||||
restored_image = np.zeros(original_shape, dtype=np.uint8)
|
||||
|
||||
# Resize face to original crop size
|
||||
resized_face = cv2.resize(processed_face_np, (w, h), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Place face in position
|
||||
restored_image[y1:y1 + h, x1:x1 + w] = resized_face
|
||||
|
||||
# Rotate back if needed
|
||||
if rotation_angle != 0:
|
||||
height, width = original_image_shape[:2]
|
||||
height, width = original_shape[:2]
|
||||
center = (width // 2, height // 2)
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, -rotation_angle, 1.0)
|
||||
restored_image = cv2.warpAffine(restored_image, rotation_matrix, (width, height),
|
||||
flags=cv2.INTER_LANCZOS4)
|
||||
restored_image = cv2.warpAffine(
|
||||
restored_image,
|
||||
rotation_matrix,
|
||||
(width, height),
|
||||
flags=cv2.INTER_LANCZOS4
|
||||
)
|
||||
|
||||
# Convert to tensor format
|
||||
# Convert to tensor
|
||||
restored_image = self.image_processor.numpy_to_tensor(restored_image)
|
||||
mask = self._create_mask(original_image_shape, crop_bbox, rotation_angle)
|
||||
|
||||
return restored_image, fp_pipe, mask
|
||||
# Create mask
|
||||
mask = self._create_mask(original_shape, (x1, y1, w, h), rotation_angle)
|
||||
|
||||
return restored_image, frame_settings, mask
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in restore: {str(e)}")
|
||||
return image, fp_pipe, self._create_empty_mask(image)
|
||||
return None
|
||||
|
||||
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
|
||||
# Create face area mask
|
||||
mask[y1:y1 + h, x1:x1 + w] = 1.0
|
||||
|
||||
# Rotate the mask back to the original orientation
|
||||
# Rotate mask if needed
|
||||
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)
|
||||
mask = cv2.warpAffine(
|
||||
mask,
|
||||
rotation_matrix,
|
||||
(width, height),
|
||||
flags=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
return self.image_processor.convert_mask_to_tensor(mask)
|
||||
|
||||
def _create_empty_mask(self, image):
|
||||
"""Create an empty mask with the same shape as the input image."""
|
||||
"""Create an empty mask matching the image dimensions."""
|
||||
if torch.is_tensor(image):
|
||||
image_shape = image.shape[1:3] # (H, W)
|
||||
shape = image.shape[1:3]
|
||||
else:
|
||||
image_shape = image.shape[:2] # (H, W)
|
||||
mask = torch.zeros((1, *image_shape), dtype=torch.float32)
|
||||
return mask
|
||||
shape = image.shape[:2]
|
||||
return torch.zeros((1, *shape), dtype=torch.float32)
|
||||
|
||||
def _load_image_from_path(self, image_path):
|
||||
"""Load and convert image from file path to tensor format."""
|
||||
|
||||
+85
-18
@@ -90,6 +90,7 @@ class FaceWrapper:
|
||||
if det_overlay is not None:
|
||||
overlays.append(det_overlay)
|
||||
|
||||
# Get base landmarks
|
||||
base_landmarks = MediapipeBaseLandmarks.get_base_landmarks(
|
||||
(width, height), x_scale=x_scale, y_translation=y_transform
|
||||
)
|
||||
@@ -117,12 +118,16 @@ class FaceWrapper:
|
||||
|
||||
output_image = self.image_processor.numpy_to_tensor(result_image)
|
||||
output_mask = self.image_processor.convert_mask_to_tensor(mask_np)
|
||||
landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks)
|
||||
|
||||
return output_image, self._update_pipe(fp_pipe, landmarks_data), output_mask
|
||||
# Prepare and update landmarks data
|
||||
landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks)
|
||||
updated_pipe = self._update_pipe(fp_pipe, landmarks_data)
|
||||
|
||||
return output_image, updated_pipe, output_mask
|
||||
|
||||
def _unwrap_mode(self, image_np, landmarks_df, width, height, device,
|
||||
x_scale, y_transform, fp_pipe, mask_np):
|
||||
# Get base landmarks
|
||||
base_landmarks = MediapipeBaseLandmarks.get_base_landmarks(
|
||||
(width, height), x_scale=x_scale, y_translation=y_transform
|
||||
)
|
||||
@@ -141,8 +146,11 @@ class FaceWrapper:
|
||||
|
||||
output_image = self.image_processor.pil_to_tensor(warped_image)
|
||||
|
||||
# Prepare and update landmarks data
|
||||
landmarks_data = self._prepare_landmarks_data(landmarks_df, base_landmarks)
|
||||
return output_image, self._update_pipe(fp_pipe, landmarks_data), warped_mask
|
||||
updated_pipe = self._update_pipe(fp_pipe, landmarks_data)
|
||||
|
||||
return output_image, updated_pipe, warped_mask
|
||||
|
||||
def _wrap_mode(self, image_np, landmarks_df, width, height, device,
|
||||
x_scale, y_transform, fp_pipe, mask_np):
|
||||
@@ -152,12 +160,25 @@ class FaceWrapper:
|
||||
output_mask = self.image_processor.convert_mask_to_tensor(mask_np)
|
||||
return output_image, {}, output_mask
|
||||
|
||||
current_frame = fp_pipe.get("current_frame", 0)
|
||||
frame_key = f"frame_{current_frame}"
|
||||
|
||||
# Check if we have required data
|
||||
if frame_key not in fp_pipe["frames"] or "detected_lm" not in fp_pipe["frames"][frame_key]:
|
||||
print(f"No detected landmarks found for frame {current_frame}")
|
||||
output_image = self.image_processor.numpy_to_tensor(image_np)
|
||||
output_mask = self.image_processor.convert_mask_to_tensor(mask_np)
|
||||
return output_image, fp_pipe, output_mask
|
||||
|
||||
# Get source landmarks from target_lm at root level
|
||||
source_x = fp_pipe['target_lm']['x']
|
||||
source_y = fp_pipe['target_lm']['y']
|
||||
source_landmarks = np.column_stack((source_x, source_y))[:468]
|
||||
|
||||
detected_x = fp_pipe['detected_lm']['x']
|
||||
detected_y = fp_pipe['detected_lm']['y']
|
||||
# Get target landmarks from detected_lm in current frame
|
||||
frame_data = fp_pipe["frames"][frame_key]
|
||||
detected_x = frame_data['detected_lm']['x']
|
||||
detected_y = frame_data['detected_lm']['y']
|
||||
target_landmarks = np.column_stack((detected_x, detected_y))[:468]
|
||||
|
||||
# Process image
|
||||
@@ -183,18 +204,64 @@ class FaceWrapper:
|
||||
return CPUDeformer.warp_face(image, source_landmarks, target_landmarks)
|
||||
|
||||
def _prepare_landmarks_data(self, detected_df, target_lm):
|
||||
return {
|
||||
'detected_lm': {
|
||||
'x': detected_df['x'].tolist(),
|
||||
'y': detected_df['y'].tolist(),
|
||||
'indices': detected_df['index'].tolist()
|
||||
},
|
||||
'target_lm': {
|
||||
'x': target_lm[:, 0].tolist(),
|
||||
'y': target_lm[:, 1].tolist(),
|
||||
'indices': list(range(len(target_lm)))
|
||||
}
|
||||
"""
|
||||
Prepare landmarks data in the correct format for fp_pipe
|
||||
|
||||
Args:
|
||||
detected_df: DataFrame with detected landmarks
|
||||
target_lm: Target landmarks array
|
||||
|
||||
Returns:
|
||||
dict: Landmarks data in the format compatible with fp_pipe
|
||||
"""
|
||||
# Prepare detected landmarks data
|
||||
detected_data = {
|
||||
'x': detected_df['x'].tolist(),
|
||||
'y': detected_df['y'].tolist(),
|
||||
'indices': detected_df['index'].tolist()
|
||||
}
|
||||
|
||||
def _update_pipe(self, pipe, new_data):
|
||||
return {**(pipe or {}), **new_data}
|
||||
# Prepare target landmarks data
|
||||
target_data = {
|
||||
'x': target_lm[:, 0].tolist(),
|
||||
'y': target_lm[:, 1].tolist(),
|
||||
'indices': list(range(len(target_lm)))
|
||||
}
|
||||
|
||||
return detected_data, target_data
|
||||
|
||||
def _update_pipe(self, pipe, landmarks_data):
|
||||
"""
|
||||
Update fp_pipe with new landmarks data
|
||||
|
||||
Args:
|
||||
pipe: Existing fp_pipe dictionary or None
|
||||
landmarks_data: Tuple of (detected_landmarks, target_landmarks)
|
||||
|
||||
Returns:
|
||||
dict: Updated fp_pipe structure
|
||||
"""
|
||||
if pipe is None:
|
||||
pipe = {
|
||||
"workflow": "single_image",
|
||||
"current_frame": 0,
|
||||
"frames": {}
|
||||
}
|
||||
|
||||
detected_lm, target_lm = landmarks_data
|
||||
|
||||
# Update target landmarks at the root level
|
||||
pipe["target_lm"] = target_lm
|
||||
|
||||
# Update detected landmarks for the current frame
|
||||
current_frame = pipe.get("current_frame", 0)
|
||||
frame_key = f"frame_{current_frame}"
|
||||
|
||||
# Ensure the frame exists in the structure
|
||||
if frame_key not in pipe["frames"]:
|
||||
pipe["frames"][frame_key] = {}
|
||||
|
||||
# Update detected landmarks for the current frame
|
||||
pipe["frames"][frame_key]["detected_lm"] = detected_lm
|
||||
|
||||
return pipe
|
||||
Reference in New Issue
Block a user