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)