Update app.py

app.py

@@ -4,31 +4,116 @@ import numpy as np
 import cv2
 import gradio as gr
 
-# 下载模型
+# Download face encoder
 snapshot_download(
     "fal/AuraFace-v1",
     local_dir="models/auraface",
 )
 
-# 初始化FaceAnalysis
-face_app = FaceAnalysis(
+app = FaceAnalysis(
     name="auraface",
     providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
     root=".",
 )
 
+app.prepare(ctx_id=0, det_size=(640, 640))
+
+def process_image_by_bbox_larger(input_image, bbox_xyxy, min_bbox_ratio=0.2):
+    """
+    Process an image based on a bounding box, cropping and resizing as necessary.
+    Parameters:
+    - input_image: PIL Image object.
+    - bbox_xyxy: Tuple (x1, y1, x2, y2) representing the bounding box coordinates.
+    Returns:
+    - A processed image cropped and resized to 1024x1024 if the bounding box is valid,
+      or None if the bounding box does not meet the required size criteria.
+    """
+    # Constants
+    target_size = 1024
+    # min_bbox_ratio = 0.2  # Bounding box should be at least 20% of the crop
+
+    # Extract bounding box coordinates
+    x1, y1, x2, y2 = bbox_xyxy
+    bbox_w = x2 - x1
+    bbox_h = y2 - y1
+
+    # Calculate the area of the bounding box
+    bbox_area = bbox_w * bbox_h
+
+    # Start with the smallest square crop that allows bbox to be at least 20% of the crop area
+    crop_size = max(bbox_w, bbox_h)
+    initial_crop_area = crop_size * crop_size
+    while (bbox_area / initial_crop_area) < min_bbox_ratio:
+        crop_size += 10  # Gradually increase until bbox is at least 20% of the area
+        initial_crop_area = crop_size * crop_size
+
+    # Once the minimum condition is satisfied, try to expand the crop further
+    max_possible_crop_size = min(input_image.width, input_image.height)
+    while crop_size < max_possible_crop_size:
+        # Calculate a potential new area
+        new_crop_size = crop_size + 10
+        new_crop_area = new_crop_size * new_crop_size
+        if (bbox_area / new_crop_area) < min_bbox_ratio:
+            break  # Stop if expanding further violates the 20% rule
+        crop_size = new_crop_size
+
+    # Determine the center of the bounding box
+    center_x = (x1 + x2) // 2
+    center_y = (y1 + y2) // 2
+
+    # Calculate the crop coordinates centered around the bounding box
+    crop_x1 = max(0, center_x - crop_size // 2)
+    crop_y1 = max(0, center_y - crop_size // 2)
+    crop_x2 = min(input_image.width, crop_x1 + crop_size)
+    crop_y2 = min(input_image.height, crop_y1 + crop_size)
+
+    # Ensure the crop is square, adjust if it goes out of image bounds
+    if crop_x2 - crop_x1 != crop_y2 - crop_y1:
+        side_length = min(crop_x2 - crop_x1, crop_y2 - crop_y1)
+        crop_x2 = crop_x1 + side_length
+        crop_y2 = crop_y1 + side_length
+
+    # Crop the image
+    cropped_image = input_image.crop((crop_x1, crop_y1, crop_x2, crop_y2))
+
+    # Resize the cropped image to 1024x1024
+    resized_image = cropped_image.resize((target_size, target_size), Image.LANCZOS)
+
+    return resized_image
+
+def calc_emb_cropped(image, app, min_bbox_ratio=0.2):
+    face_image = image.copy()
+
+    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
+
+    face_info = face_info[0]
+
+    #print(face_info)
+
+    cropped_face_image = process_image_by_bbox_larger(face_image, face_info["bbox"], min_bbox_ratio=min_bbox_ratio)
+
+    return cropped_face_image
+
 def get_embedding(image):
-    # 将图片转换为OpenCV格式
-    cv2_image = np.array(image.convert("RGB"))
-    cv2_image = cv2_image[:, :, ::-1]
+    face_image = image.copy()
+
+    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
     
     # 获取人脸嵌入
-    faces = face_app.get(cv2_image)
-    if len(faces) > 0:
-        return faces[0].normed_embedding
+    #face_info = app.get(cv2_image)
+    if len(face_info) > 0:
+        return face_info[0].normed_embedding
     else:
         return None
 
+'''
+from PIL import Image 
+im0 = Image.open("Unreal_5_render_of_a_handsome_man_gentle_snowfall_at_dusk_a_bustling_marketplace_in_the_background.png")
+calc_emb_cropped(im0, app)
+get_embedding(im0)
+'''
+
+
 def calculate_similarity(image1, image2):
     # 获取两张图片的嵌入
     embedding1 = get_embedding(image1)