Hugging Face's logo

Spaces:
svjack
/
auraface-embedding
Sleeping

App Files Community
auraface-embedding
File size: 2,995 Bytes 
00e256c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
 
import os
from datasets import load_dataset
import uuid
import numpy as np
import cv2
import gradio as gr
from huggingface_hub import snapshot_download
from insightface.app import FaceAnalysis
from PIL import Image
import json

# 定义保存路径
save_path = "./examples/xiangxiang_man"

# 清空目标路径（如果存在）
if os.path.exists(save_path):
    for file_name in os.listdir(save_path):
        file_path = os.path.join(save_path, file_name)
        if os.path.isfile(file_path):
            os.remove(file_path)
    print(f"Cleared existing files in {save_path}")
else:
    os.makedirs(save_path, exist_ok=True)
    print(f"Created directory: {save_path}")

# 加载数据集
dataset = load_dataset("svjack/Prince_Xiang_iclight_v2")

# 遍历数据集并保存图片
for example in dataset["train"]:
    # 获取图片数据
    image = example["image"]
    
    # 生成唯一的文件名（使用 uuid）
    file_name = f"{uuid.uuid4()}.png"
    file_path = os.path.join(save_path, file_name)
    
    # 保存图片
    image.save(file_path)
    print(f"Saved {file_path}")

print("All images have been saved.")

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

# Initialize FaceAnalysis
app = FaceAnalysis(
    name="auraface",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
    root=".",
)
app.prepare(ctx_id=0, det_size=(640, 640))

def get_embedding(image):
    """
    Get the embedding of a single image.
    Parameters:
    - image: PIL Image object.
    Returns:
    - A numpy array representing the embedding of the face in the image.
    """
    # Convert PIL image to OpenCV format
    cv2_image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # Get face information
    face_info = app.get(cv2_image)
    
    if len(face_info) > 0:
        # Return the embedding of the first detected face
        return face_info[0].normed_embedding.tolist()  # Convert to list
    else:
        return None

def display_embedding(image):
    """
    Display the embedding of a single image as a JSON object.
    Parameters:
    - image: PIL Image object.
    Returns:
    - A JSON object with the embedding (nested list) or an empty list if no face is detected.
    """
    embedding = get_embedding(image)
    if embedding is not None:
        return json.dumps({"embedding": embedding})  # Wrap in a list and convert to JSON
    else:
        return json.dumps({"embedding": []})  # Return empty list as JSON

# 获取数据集中的图片路径
import pathlib
example_images = list(map(str, pathlib.Path(save_path).rglob("*.png")))

# 创建Gradio界面
iface = gr.Interface(
    fn=display_embedding,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="面部图片嵌入计算",
    description="上传一张图片，计算其嵌入向量。",
    examples=example_images[:3],  # 使用数据集中的前3张图片作为示例
)

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