476 lines
16 KiB
Python
476 lines
16 KiB
Python
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 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 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)
|
|
|
|
@staticmethod
|
|
def draw_landmarks(image_size: Union[int, Tuple[int, int]],
|
|
landmarks_df: pd.DataFrame,
|
|
transparency: float = 0.5,
|
|
color: Tuple[int, int, int] = (0, 255, 0),
|
|
radius: int = 2,
|
|
label: bool = False) -> Union[np.ndarray, None]:
|
|
"""
|
|
Draw facial landmarks on a transparent image.
|
|
|
|
Args:
|
|
image_size: Target image size (int or tuple of width, height)
|
|
landmarks_df: DataFrame containing landmark coordinates
|
|
transparency: Alpha channel value (0.0-1.0)
|
|
color: RGB color tuple for landmarks
|
|
radius: Radius of landmark points
|
|
label: Whether to show landmark IDs
|
|
|
|
Returns:
|
|
numpy array: RGBA image with drawn landmarks
|
|
"""
|
|
if landmarks_df is None or landmarks_df.empty:
|
|
return None
|
|
|
|
# Create transparent image
|
|
if isinstance(image_size, tuple):
|
|
width, height = image_size
|
|
else:
|
|
width = height = image_size
|
|
|
|
# Create RGBA image (alpha channel for transparency)
|
|
image = np.zeros((height, width, 4), dtype=np.uint8)
|
|
|
|
# Set alpha channel based on transparency
|
|
alpha = int(transparency * 255)
|
|
|
|
# Draw landmarks
|
|
for _, landmark in landmarks_df.iterrows():
|
|
x = int(landmark['x'])
|
|
y = int(landmark['y'])
|
|
index = int(landmark['index'])
|
|
|
|
# Ensure coordinates are within image bounds
|
|
if 0 <= x < width and 0 <= y < height:
|
|
# Draw filled circle
|
|
cv2.circle(
|
|
image,
|
|
(x, y),
|
|
radius,
|
|
(*color, alpha),
|
|
-1, # Filled circle
|
|
cv2.LINE_AA
|
|
)
|
|
|
|
# Add landmark ID if requested
|
|
if label:
|
|
# Text parameters
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
font_scale = 0.4
|
|
text_thickness = 1
|
|
|
|
# Add small offset to text position
|
|
text_x = x + radius + 5
|
|
text_y = y + radius
|
|
|
|
# Draw text with white color and black outline for better visibility
|
|
text = str(index)
|
|
|
|
# Draw text outline
|
|
cv2.putText(
|
|
image,
|
|
text,
|
|
(text_x, text_y),
|
|
font,
|
|
font_scale,
|
|
(0, 0, 0, alpha), # Black outline
|
|
text_thickness + 1,
|
|
cv2.LINE_AA
|
|
)
|
|
|
|
# Draw text
|
|
cv2.putText(
|
|
image,
|
|
text,
|
|
(text_x, text_y),
|
|
font,
|
|
font_scale,
|
|
(*color, alpha), # Main color
|
|
text_thickness,
|
|
cv2.LINE_AA
|
|
)
|
|
|
|
return image
|
|
|
|
@staticmethod
|
|
def draw_dynamic_histogram(composite: np.ndarray,
|
|
input_black: int, input_white: int,
|
|
gamma: float) -> np.ndarray:
|
|
"""Creates histogram visualization with dynamic control lines"""
|
|
# Convert to grayscale for histogram calculation
|
|
gray = cv2.cvtColor(composite, cv2.COLOR_RGB2GRAY) if len(composite.shape) == 3 else composite
|
|
|
|
# Calculate histogram
|
|
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
|
|
hist_img = np.zeros((200, 256, 3), dtype=np.uint8)
|
|
cv2.normalize(hist, hist, 0, hist_img.shape[0], cv2.NORM_MINMAX)
|
|
|
|
# Draw histogram baseline
|
|
for i in range(1, 256):
|
|
cv2.line(hist_img,
|
|
(i - 1, 200 - int(hist[i - 1])),
|
|
(i, 200 - int(hist[i])),
|
|
(128, 128, 128), 1)
|
|
|
|
# Calculate gamma position
|
|
gamma_pos = int(255 * (0.5 ** (1 / gamma)))
|
|
|
|
# Draw control lines
|
|
cv2.line(hist_img, (input_black, 0), (input_black, 200), (255, 0, 0), 2) # Blue - black level
|
|
cv2.line(hist_img, (input_white, 0), (input_white, 200), (0, 255, 0), 2) # Green - white level
|
|
cv2.line(hist_img, (gamma_pos, 0), (gamma_pos, 200), (255, 255, 255), 2) # White - gamma
|
|
|
|
# Add legend
|
|
cv2.putText(hist_img, f"Black: {input_black}", (10, 20),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
|
|
cv2.putText(hist_img, f"White: {input_white}", (10, 40),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
|
|
cv2.putText(hist_img, f"Gamma: {gamma:.2f}", (10, 60),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
|
|
|
|
return hist_img
|
|
|
|
# Image Converters
|
|
@staticmethod
|
|
def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray:
|
|
"""Convert PyTorch tensor to numpy array.
|
|
|
|
Args:
|
|
tensor: Input tensor in format (B,H,W,C) or (B,C,H,W), values 0-1
|
|
|
|
Returns:
|
|
np.ndarray: RGB image (H,W,3), uint8 values 0-255
|
|
"""
|
|
if tensor is None:
|
|
return None
|
|
|
|
# Handle different tensor formats
|
|
img = tensor.detach().cpu().numpy()
|
|
if len(img.shape) == 4:
|
|
img = img[0] # Remove batch dimension
|
|
|
|
# Handle channel dimension if in (C,H,W) format
|
|
if img.shape[0] in [1, 3, 4]:
|
|
img = img.transpose(1, 2, 0)
|
|
|
|
# Handle single channel
|
|
if len(img.shape) == 2 or img.shape[2] == 1:
|
|
img = np.stack((img,) * 3, axis=-1)
|
|
|
|
# Handle RGBA
|
|
elif img.shape[2] == 4:
|
|
img = img[..., :3]
|
|
|
|
# Convert to uint8 if float
|
|
if img.dtype == np.float32 or img.dtype == np.float64:
|
|
img = (img * 255).clip(0, 255).astype(np.uint8)
|
|
|
|
return img
|
|
|
|
@staticmethod
|
|
def numpy_to_tensor(array: np.ndarray) -> torch.Tensor:
|
|
"""Convert numpy array to PyTorch tensor.
|
|
|
|
Args:
|
|
array: RGB image (H,W,3), uint8 or float32
|
|
|
|
Returns:
|
|
torch.Tensor: (1,H,W,3) tensor, values 0-1
|
|
"""
|
|
if array is None:
|
|
return None
|
|
|
|
# Convert to float32 if uint8
|
|
if array.dtype == np.uint8:
|
|
array = array.astype(np.float32) / 255.0
|
|
|
|
# Add batch dimension
|
|
tensor = torch.from_numpy(array).unsqueeze(0)
|
|
return tensor
|
|
|
|
@staticmethod
|
|
def pil_to_numpy(image: Image.Image) -> np.ndarray:
|
|
"""Convert PIL image to numpy array.
|
|
|
|
Args:
|
|
image: PIL Image
|
|
|
|
Returns:
|
|
np.ndarray: RGB image (H,W,3), uint8 values 0-255
|
|
"""
|
|
if image is None:
|
|
return None
|
|
|
|
# Convert to RGB if needed
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
|
|
return np.array(image)
|
|
|
|
@staticmethod
|
|
def numpy_to_pil(array: np.ndarray) -> Image.Image:
|
|
"""Convert numpy array to PIL image.
|
|
|
|
Args:
|
|
array: RGB image (H,W,3), uint8 or float32
|
|
|
|
Returns:
|
|
PIL.Image: RGB PIL image
|
|
"""
|
|
if array is None:
|
|
return None
|
|
|
|
# Convert to uint8 if float
|
|
if array.dtype == np.float32 or array.dtype == np.float64:
|
|
array = (array * 255).clip(0, 255).astype(np.uint8)
|
|
|
|
return Image.fromarray(array)
|
|
|
|
@staticmethod
|
|
def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
|
"""Convert PyTorch tensor to PIL image.
|
|
|
|
Args:
|
|
tensor: Input tensor in format (B,H,W,C), values 0-1
|
|
|
|
Returns:
|
|
PIL.Image: RGB PIL image
|
|
"""
|
|
return ImageProcessor.numpy_to_pil(ImageProcessor.tensor_to_numpy(tensor))
|
|
|
|
@staticmethod
|
|
def pil_to_tensor(image: Image.Image) -> torch.Tensor:
|
|
"""Convert PIL image to PyTorch tensor.
|
|
|
|
Args:
|
|
image: PIL Image
|
|
|
|
Returns:
|
|
torch.Tensor: (1,H,W,3) tensor, values 0-1
|
|
"""
|
|
return ImageProcessor.numpy_to_tensor(ImageProcessor.pil_to_numpy(image))
|
|
|
|
@staticmethod
|
|
def convert_to_numpy(image: Union[torch.Tensor, np.ndarray, Image.Image]) -> np.ndarray:
|
|
"""Universal converter to numpy array.
|
|
|
|
Args:
|
|
image: Input in any supported format
|
|
|
|
Returns:
|
|
np.ndarray: RGB image (H,W,3), uint8 values 0-255
|
|
"""
|
|
if image is None:
|
|
return None
|
|
|
|
if torch.is_tensor(image):
|
|
return ImageProcessor.tensor_to_numpy(image)
|
|
elif isinstance(image, Image.Image):
|
|
return ImageProcessor.pil_to_numpy(image)
|
|
elif isinstance(image, np.ndarray):
|
|
# Handle grayscale
|
|
if len(image.shape) == 2:
|
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
|
# Handle RGBA
|
|
elif image.shape[2] == 4:
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
|
|
# Convert to uint8 if float
|
|
if image.dtype == np.float32 or image.dtype == np.float64:
|
|
image = (image * 255).clip(0, 255).astype(np.uint8)
|
|
return image
|
|
else:
|
|
raise ValueError(f"Unsupported image type: {type(image)}")
|
|
|
|
# Mask conversion methods
|
|
@staticmethod
|
|
def convert_mask_to_numpy(mask: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
|
|
"""Convert mask to numpy array.
|
|
|
|
Args:
|
|
mask: Input mask as tensor (1,H,W) or array
|
|
|
|
Returns:
|
|
np.ndarray: Mask as (H,W) uint8 array, values 0-255
|
|
"""
|
|
if mask is None:
|
|
return None
|
|
|
|
if torch.is_tensor(mask):
|
|
mask = mask.detach().cpu().numpy()
|
|
if len(mask.shape) == 3:
|
|
mask = mask[0] # Remove batch dimension
|
|
|
|
if mask.dtype == np.float32 or mask.dtype == np.float64:
|
|
mask = (mask * 255).clip(0, 255).astype(np.uint8)
|
|
|
|
return mask
|
|
|
|
@staticmethod
|
|
def convert_mask_to_tensor(mask: np.ndarray) -> torch.Tensor:
|
|
"""Convert numpy mask to tensor.
|
|
|
|
Args:
|
|
mask: Input mask as (H,W) array
|
|
|
|
Returns:
|
|
torch.Tensor: Mask as (1,H,W) tensor, values 0-1
|
|
"""
|
|
if mask is None:
|
|
return None
|
|
|
|
if mask.dtype == np.uint8:
|
|
mask = mask.astype(np.float32) / 255.0
|
|
|
|
return torch.from_numpy(mask).unsqueeze(0)
|
|
|
|
# Loaders
|
|
@staticmethod
|
|
def load_image_from_path(image_path):
|
|
"""Load and convert image from file path to tensor format."""
|
|
try:
|
|
image = Image.open(image_path)
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
return ImageProcessor.pil_to_tensor(image)
|
|
except Exception as e:
|
|
print(f"Error loading image from {image_path}: {str(e)}")
|
|
return None
|
|
|