Create app.py

app.py

@@ -0,0 +1,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)