File size: 1,670 Bytes
6ee7e4d |
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 |
# 导入所需库
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")
|