# 导入所需库
import os
from PIL import Image
import pandas as pd
from transformers import pipeline

# 任务 2: 加载预训练模型
# 加载年龄分类模型
age_classifier = pipeline("image-classification", model="nateraw/vit-age-classifier")
# 加载性别分类模型
gender_classifier = pipeline("image-classification", model="rizvandwiki/gender-classification")
# 加载幸福度分类模型
emotion_classifier = pipeline("image-classification", model="Rajaram1996/Happiness-Classifier")

# 任务 1: 定义图像文件夹路径
image_folder = "images/"
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.png'))]

# 初始化结果列表
results = []

# 任务 3: 处理每张图像并分类属性
for image_file in image_files:
    image_path = os.path.join(image_folder, image_file)
    image = Image.open(image_path)

    # 预测年龄
    age_prediction = age_classifier(image)
    predicted_age = age_prediction[0]['label']

    # 预测性别
    gender_prediction = gender_classifier(image)
    predicted_gender = gender_prediction[0]['label']

    # 预测幸福度
    emotion_prediction = emotion_classifier(image)
    predicted_happiness = emotion_prediction[0]['label']

    # 将结果添加到列表
    results.append({
        "Image Name": image_file,
        "Predicted Age": predicted_age,
        "Predicted Gender": predicted_gender,
        "Predicted Happiness": predicted_happiness
    })

# 任务 4: 使用 pandas 生成 CSV 报告
df = pd.DataFrame(results)
df.to_csv("participant_classification_report.csv", index=False)

print("分类完成，结果已保存到 participant_classification_report.csv")