Hugging Face's logo

Spaces:
svjack
/
auraface-similarity
Running

App Files Community
auraface-similarity
File size: 4,630 Bytes 
6d656a7
 
 
 
 
 
cca4277
6d656a7
 
 
 
 
cca4277
6d656a7
 
 
 
 
cca4277
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6d656a7
cca4277
 
 
6d656a7
 
cca4277
 
 
6d656a7
 
 
cca4277
 
 
 
 
 
 
 
6d656a7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
 
from huggingface_hub import snapshot_download
from insightface.app import FaceAnalysis
import numpy as np
import cv2
import gradio as gr

# Download face encoder
snapshot_download(
    "fal/AuraFace-v1",
    local_dir="models/auraface",
)

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):
    face_image = image.copy()

    face_info = app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR))
    
    # 获取人脸嵌入
    #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)
    embedding2 = get_embedding(image2)
    
    if embedding1 is not None and embedding2 is not None:
        # 计算余弦相似度
        similarity = np.dot(embedding1, embedding2) / (np.linalg.norm(embedding1) * np.linalg.norm(embedding2))
        return f"图片相似度: {similarity:.4f}"
    else:
        return "无法检测到人脸或计算相似度"

# 创建Gradio界面
iface = gr.Interface(
    fn=calculate_similarity,
    inputs=[gr.Image(type="pil"), gr.Image(type="pil")],
    outputs="text",
    title="图片相似度计算",
    description="上传两张图片，计算它们的相似度。"
)

# 启动Gradio应用
iface.launch(share = True)