412 lines
15 KiB
Python
412 lines
15 KiB
Python
# Copyright (c) 2025 Fudan University. All rights reserved.
|
|
|
|
|
|
from io import BytesIO
|
|
import random
|
|
from PIL import Image
|
|
import numpy as np
|
|
import cv2
|
|
import insightface
|
|
import torch
|
|
from torchvision import transforms
|
|
from torch.cuda.amp import autocast
|
|
|
|
def face_preserving_resize(img, face_bboxes, target_size=512):
|
|
"""
|
|
Resize image while ensuring all faces are preserved in the output.
|
|
|
|
Args:
|
|
img: PIL Image
|
|
face_bboxes: List of [x1, y1, x2, y2] face coordinates
|
|
target_size: Maximum dimension for resizing
|
|
|
|
Returns:
|
|
Tuple of (resized image, new_bboxes) or (None, None) if faces can't fit
|
|
"""
|
|
|
|
x1_1, y1_1, x2_1, y2_1 = map(int, face_bboxes[0])
|
|
x1_2, y1_2, x2_2, y2_2 = map(int, face_bboxes[1])
|
|
min_x1 = min(x1_1, x1_2)
|
|
min_y1 = min(y1_1, y1_2)
|
|
max_x2 = max(x2_1, x2_2)
|
|
max_y2 = max(y2_1, y2_2)
|
|
# print("min_x1:", min_x1, "min_y1:", min_y1, "max_x2:", max_x2, "max_y2:", max_y2)
|
|
# if any of them is negative, we cannot resize (Idk why this happens)
|
|
if min_x1 < 0 or min_y1 < 0 or max_x2 < 0 or max_y2 < 0:
|
|
return None, None
|
|
|
|
# if face width is longer than the image height, or the face height is longer than the image width, we cannot resize
|
|
face_width = max_x2 - min_x1
|
|
face_height = max_y2 - min_y1
|
|
if face_width > img.height or face_height > img.width:
|
|
return None, None
|
|
|
|
# Create a copy of face_bboxes for transformation
|
|
new_bboxes = []
|
|
for bbox in face_bboxes:
|
|
new_bboxes.append(list(map(int, bbox)))
|
|
|
|
# Choose cropping strategy based on image aspect ratio
|
|
if img.width > img.height:
|
|
# We need to crop width to make a square
|
|
square_size = img.height
|
|
|
|
# Calculate valid horizontal crop range that preserves all faces
|
|
left_max = min_x1 # Leftmost position that includes leftmost face
|
|
right_min = max_x2 - square_size # Rightmost position that includes rightmost face
|
|
|
|
if right_min <= left_max:
|
|
# We can find a valid crop window
|
|
start = random.randint(int(right_min), int(left_max)) if right_min < left_max else int(right_min)
|
|
start = max(0, min(start, img.width - square_size)) # Ensure within image bounds
|
|
else:
|
|
# Faces are too far apart for square crop - use center of faces
|
|
face_center = (min_x1 + max_x2) // 2
|
|
start = max(0, min(face_center - (square_size // 2), img.width - square_size))
|
|
|
|
cropped_img = img.crop((start, 0, start + square_size, square_size))
|
|
|
|
# Adjust bounding box coordinates based on crop
|
|
for bbox in new_bboxes:
|
|
bbox[0] -= start # x1 adjustment
|
|
bbox[2] -= start # x2 adjustment
|
|
# y coordinates remain unchanged
|
|
else:
|
|
# We need to crop height to make a square
|
|
square_size = img.width
|
|
|
|
# Calculate valid vertical crop range that preserves all faces
|
|
top_max = min_y1 # Topmost position that includes topmost face
|
|
bottom_min = max_y2 - square_size # Bottommost position that includes bottommost face
|
|
|
|
if bottom_min <= top_max:
|
|
# We can find a valid crop window
|
|
start = random.randint(int(bottom_min), int(top_max)) if bottom_min < top_max else int(bottom_min)
|
|
start = max(0, min(start, img.height - square_size)) # Ensure within image bounds
|
|
else:
|
|
# Faces are too far apart for square crop - use center of faces
|
|
face_center = (min_y1 + max_y2) // 2
|
|
start = max(0, min(face_center - (square_size // 2), img.height - square_size))
|
|
|
|
cropped_img = img.crop((0, start, square_size, start + square_size))
|
|
|
|
# Adjust bounding box coordinates based on crop
|
|
for bbox in new_bboxes:
|
|
bbox[1] -= start # y1 adjustment
|
|
bbox[3] -= start # y2 adjustment
|
|
# x coordinates remain unchanged
|
|
|
|
# Calculate scale factor for resizing from square_size to target_size
|
|
scale_factor = target_size / square_size
|
|
|
|
# Adjust bounding boxes for the resize operation
|
|
for bbox in new_bboxes:
|
|
bbox[0] = int(bbox[0] * scale_factor)
|
|
bbox[1] = int(bbox[1] * scale_factor)
|
|
bbox[2] = int(bbox[2] * scale_factor)
|
|
bbox[3] = int(bbox[3] * scale_factor)
|
|
|
|
# Final resize to target size
|
|
resized_img = cropped_img.resize((target_size, target_size), Image.Resampling.LANCZOS)
|
|
|
|
# Make sure all coordinates are within bounds (0 to target_size)
|
|
# for bbox in new_bboxes:
|
|
# bbox[0] = max(0, min(bbox[0], target_size - 1))
|
|
# bbox[1] = max(0, min(bbox[1], target_size - 1))
|
|
# bbox[2] = max(1, min(bbox[2], target_size))
|
|
# bbox[3] = max(1, min(bbox[3], target_size))
|
|
|
|
return resized_img, new_bboxes
|
|
|
|
def extract_moref(img, json_data, face_size_restriction=100):
|
|
"""
|
|
Extract faces from an image based on bounding boxes in JSON data.
|
|
Makes each face square and resizes to 512x512.
|
|
|
|
Args:
|
|
img: PIL Image or image data
|
|
json_data: JSON object with 'bboxes' and 'crop' information
|
|
|
|
Returns:
|
|
List of PIL Images, each 512x512, containing extracted faces
|
|
"""
|
|
# Ensure img is a PIL Image
|
|
try:
|
|
if not isinstance(img, Image.Image) and not isinstance(img, torch.Tensor) and not isinstance(img, JpegImageFile):
|
|
img = Image.open(BytesIO(img))
|
|
|
|
bboxes = json_data['bboxes']
|
|
# crop = json_data['crop']
|
|
# print("len of bboxes:", len(bboxes))
|
|
# Recalculate bounding boxes based on crop info
|
|
# new_bboxes = [recalculate_bbox(bbox, crop) for bbox in bboxes]
|
|
new_bboxes = bboxes
|
|
# any of the face is less than 100 * 100, we ignore this image
|
|
for bbox in new_bboxes:
|
|
x1, y1, x2, y2 = bbox
|
|
if x2 - x1 < face_size_restriction or y2 - y1 < face_size_restriction:
|
|
return []
|
|
# print("len of new_bboxes:", len(new_bboxes))
|
|
faces = []
|
|
for bbox in new_bboxes:
|
|
# print("processing bbox")
|
|
# Convert coordinates to integers
|
|
x1, y1, x2, y2 = map(int, bbox)
|
|
|
|
# Calculate width and height
|
|
width = x2 - x1
|
|
height = y2 - y1
|
|
|
|
# Make the bounding box square by expanding the shorter dimension
|
|
if width > height:
|
|
# Height is shorter, expand it
|
|
diff = width - height
|
|
y1 -= diff // 2
|
|
y2 += diff - (diff // 2) # Handle odd differences
|
|
elif height > width:
|
|
# Width is shorter, expand it
|
|
diff = height - width
|
|
x1 -= diff // 2
|
|
x2 += diff - (diff // 2) # Handle odd differences
|
|
|
|
# Ensure coordinates are within image boundaries
|
|
img_width, img_height = img.size
|
|
x1 = max(0, x1)
|
|
y1 = max(0, y1)
|
|
x2 = min(img_width, x2)
|
|
y2 = min(img_height, y2)
|
|
|
|
# Extract face region
|
|
face_region = img.crop((x1, y1, x2, y2))
|
|
|
|
# Resize to 512x512
|
|
face_region = face_region.resize((512, 512), Image.LANCZOS)
|
|
|
|
faces.append(face_region)
|
|
# print("len of faces:", len(faces))
|
|
return faces
|
|
except Exception as e:
|
|
print(f"Error processing image: {e}")
|
|
return []
|
|
|
|
def general_face_preserving_resize(img, face_bboxes, target_size=512):
|
|
"""
|
|
Resize image while ensuring all faces are preserved in the output.
|
|
Handles any number of faces (1-5).
|
|
|
|
Args:
|
|
img: PIL Image
|
|
face_bboxes: List of [x1, y1, x2, y2] face coordinates
|
|
target_size: Maximum dimension for resizing
|
|
|
|
Returns:
|
|
Tuple of (resized image, new_bboxes) or (None, None) if faces can't fit
|
|
"""
|
|
# Find bounding region containing all faces
|
|
if not face_bboxes:
|
|
print("Warning: No face bounding boxes provided.")
|
|
return None, None
|
|
|
|
min_x1 = min(bbox[0] for bbox in face_bboxes)
|
|
min_y1 = min(bbox[1] for bbox in face_bboxes)
|
|
max_x2 = max(bbox[2] for bbox in face_bboxes)
|
|
max_y2 = max(bbox[3] for bbox in face_bboxes)
|
|
|
|
# Check for negative coordinates
|
|
if min_x1 < 0 or min_y1 < 0 or max_x2 < 0 or max_y2 < 0:
|
|
# print("Warning: Negative coordinates found in face bounding boxes.")
|
|
# return None, None
|
|
min_x1 = max(min_x1, 0)
|
|
min_y1 = max(min_y1, 0)
|
|
|
|
# Check if faces fit within image
|
|
face_width = max_x2 - min_x1
|
|
face_height = max_y2 - min_y1
|
|
if face_width > img.height or face_height > img.width:
|
|
# print("Warning: Faces are too large for the image dimensions.")
|
|
# return None, None
|
|
# Instead of returning None, we will crop the image to fit the faces
|
|
max_x2 = min(max_x2, img.width)
|
|
max_y2 = min(max_y2, img.height)
|
|
min_x1 = max(min_x1, 0)
|
|
min_y1 = max(min_y1, 0)
|
|
# Create a copy of face_bboxes for transformation
|
|
new_bboxes = []
|
|
for bbox in face_bboxes:
|
|
new_bboxes.append(list(map(int, bbox)))
|
|
|
|
# Choose cropping strategy based on image aspect ratio
|
|
if img.width > img.height:
|
|
# Crop width to make a square
|
|
square_size = img.height
|
|
|
|
# Calculate valid horizontal crop range
|
|
left_max = min_x1
|
|
right_min = max_x2 - square_size
|
|
|
|
if right_min <= left_max:
|
|
# We can find a valid crop window
|
|
start = random.randint(int(right_min), int(left_max)) if right_min < left_max else int(right_min)
|
|
start = max(0, min(start, img.width - square_size))
|
|
else:
|
|
# Faces are too far apart - use center of faces
|
|
face_center = (min_x1 + max_x2) // 2
|
|
start = max(0, min(face_center - (square_size // 2), img.width - square_size))
|
|
|
|
cropped_img = img.crop((start, 0, start + square_size, square_size))
|
|
|
|
# Adjust bounding box coordinates
|
|
for bbox in new_bboxes:
|
|
bbox[0] -= start
|
|
bbox[2] -= start
|
|
else:
|
|
# Crop height to make a square
|
|
square_size = img.width
|
|
|
|
# Calculate valid vertical crop range
|
|
top_max = min_y1
|
|
bottom_min = max_y2 - square_size
|
|
|
|
if bottom_min <= top_max:
|
|
start = random.randint(int(bottom_min), int(top_max)) if bottom_min < top_max else int(bottom_min)
|
|
start = max(0, min(start, img.height - square_size))
|
|
else:
|
|
face_center = (min_y1 + max_y2) // 2
|
|
start = max(0, min(face_center - (square_size // 2), img.height - square_size))
|
|
|
|
cropped_img = img.crop((0, start, square_size, start + square_size))
|
|
|
|
# Adjust bounding box coordinates
|
|
for bbox in new_bboxes:
|
|
bbox[1] -= start
|
|
bbox[3] -= start
|
|
|
|
# Calculate scale factor and adjust bounding boxes
|
|
scale_factor = target_size / square_size
|
|
|
|
for bbox in new_bboxes:
|
|
bbox[0] = int(bbox[0] * scale_factor)
|
|
bbox[1] = int(bbox[1] * scale_factor)
|
|
bbox[2] = int(bbox[2] * scale_factor)
|
|
bbox[3] = int(bbox[3] * scale_factor)
|
|
|
|
# Final resize to target size
|
|
resized_img = cropped_img.resize((target_size, target_size), Image.Resampling.LANCZOS)
|
|
|
|
# Make sure all coordinates are within bounds
|
|
for bbox in new_bboxes:
|
|
bbox[0] = max(0, min(bbox[0], target_size - 1))
|
|
bbox[1] = max(0, min(bbox[1], target_size - 1))
|
|
bbox[2] = max(1, min(bbox[2], target_size))
|
|
bbox[3] = max(1, min(bbox[3], target_size))
|
|
|
|
return resized_img, new_bboxes
|
|
|
|
|
|
|
|
def horizontal_concat(images):
|
|
widths, heights = zip(*(img.size for img in images))
|
|
|
|
total_width = sum(widths)
|
|
max_height = max(heights)
|
|
|
|
new_im = Image.new('RGB', (total_width, max_height))
|
|
|
|
x_offset = 0
|
|
for img in images:
|
|
new_im.paste(img, (x_offset, 0))
|
|
x_offset += img.size[0]
|
|
|
|
return new_im
|
|
|
|
def extract_object(birefnet, image):
|
|
|
|
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
input_images = transforms.ToTensor()(image).unsqueeze(0).to('cuda', dtype=torch.bfloat16)
|
|
|
|
# Prediction
|
|
with torch.no_grad(), autocast(dtype=torch.bfloat16):
|
|
preds = birefnet(input_images)[-1].sigmoid().cpu()
|
|
pred = preds[0].squeeze().float()
|
|
pred_pil = transforms.ToPILImage()(pred)
|
|
mask = pred_pil.resize(image.size)
|
|
|
|
# Create a binary mask (0 or 255)
|
|
binary_mask = mask.convert("L")
|
|
|
|
# Create a new image with black background
|
|
result = Image.new("RGB", image.size, (0, 0, 0))
|
|
|
|
# Paste the original image onto the black background using the mask
|
|
result.paste(image, (0, 0), binary_mask)
|
|
|
|
return result, mask
|
|
|
|
class FaceExtractor:
|
|
def __init__(self, model_path="./"):
|
|
self.model = insightface.app.FaceAnalysis(name = "antelopev2", root=model_path)
|
|
self.model.prepare(ctx_id=0, det_thresh=0.4)
|
|
|
|
def extract(self, image: Image.Image):
|
|
image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
|
res = self.model.get(image_np)
|
|
if len(res) == 0:
|
|
return None, None
|
|
res = res[0]
|
|
# print(res.keys())
|
|
bbox = res["bbox"]
|
|
# print("len(bbox)", len(bbox))
|
|
|
|
moref = extract_moref(image, {"bboxes": [bbox]}, 1)
|
|
# print("len(moref)", len(moref))
|
|
return moref[0], res["embedding"]
|
|
|
|
def locate_bboxes(self, image: Image.Image):
|
|
image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
|
res = self.model.get(image_np)
|
|
if len(res) == 0:
|
|
return None
|
|
bboxes = []
|
|
for r in res:
|
|
bbox = r["bbox"]
|
|
bboxes.append(bbox)
|
|
|
|
_, new_bboxes_ = general_face_preserving_resize(image, bboxes, 512)
|
|
|
|
# ensure the bbox is square
|
|
new_bboxes = []
|
|
for bbox in new_bboxes_:
|
|
x1, y1, x2, y2 = bbox
|
|
w = x2 - x1
|
|
h = y2 - y1
|
|
if w > h:
|
|
diff = w - h
|
|
y1 = max(0, y1 - diff // 2)
|
|
y2 = min(512, y2 + diff // 2 + diff % 2)
|
|
else:
|
|
diff = h - w
|
|
x1 = max(0, x1 - diff // 2)
|
|
x2 = min(512, x2 + diff // 2 + diff % 2)
|
|
new_bboxes.append([x1, y1, x2, y2])
|
|
|
|
return new_bboxes
|
|
def extract_refs(self, image: Image.Image):
|
|
"""
|
|
Extracts reference faces from the image.
|
|
Returns a list of reference images and their arcface embeddings.
|
|
"""
|
|
image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
|
res = self.model.get(image_np)
|
|
if len(res) == 0:
|
|
return None, None
|
|
ref_imgs = []
|
|
arcface_embeddings = []
|
|
for r in res:
|
|
bbox = r["bbox"]
|
|
moref = extract_moref(image, {"bboxes": [bbox]}, 1)
|
|
ref_imgs.append(moref[0])
|
|
arcface_embeddings.append(r["embedding"])
|
|
return ref_imgs, arcface_embeddings
|