文章目录
- 前提
- 模型下载
- 表情分类:emotion.py
- 结果
前提
注意,输入的图片是裁剪好的人脸图,如:
检测人脸与检测参考:YOLOv7-face人脸检测
模型下载
https://download.csdn.net/download/WhiffeYF/89654401
解压后使用 emotion.pth
表情分类:emotion.py
emotion.py
执行:
python emotion.py --weights emotion.pth --source crops/face/
'''
python emotion.py --weights emotion.pth --source crops/face/
'''
import os
import torch
from torchvision import transforms
from PIL import Image
import torch.nn as nn
import torchvision.models as models
import argparseclass EmotionClassifier:def __init__(self, model_path, device='cpu'):self.device = deviceself.class_names = ['angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise']self.model = self.load_model(model_path)self.transform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])def load_model(self, model_path):model = models.mobilenet_v2(pretrained=False)model.classifier[1] = nn.Linear(model.last_channel, len(self.class_names))model.load_state_dict(torch.load(model_path, map_location=self.device))model.eval()model.to(self.device)return modeldef predict(self, image_path):image = Image.open(image_path).convert('RGB')image = self.transform(image).unsqueeze(0).to(self.device)with torch.no_grad():outputs = self.model(image)_, predicted = torch.max(outputs, 1)predicted_class = self.class_names[predicted.item()]return predicted_classdef predict_folder(self, folder_path):results = {}for filename in os.listdir(folder_path):if filename.endswith(('.png', '.jpg', '.jpeg')):file_path = os.path.join(folder_path, filename)prediction = self.predict(file_path)results[filename] = predictionreturn resultsdef classify(self, input_path):if os.path.isdir(input_path):return self.predict_folder(input_path)elif os.path.isfile(input_path):return {os.path.basename(input_path): self.predict(input_path)}else:raise ValueError(f"Invalid path: {input_path}")if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument('--weights', type=str, default='emotion.pth')parser.add_argument('--source', type=str, default='crops/face/')opt = parser.parse_args()weights = opt.weightssource = opt.source # 可以是单张图片路径或图片文件夹路径'''model_path = 'emotion.pth'input_path = 'crops/face/' '''classifier = EmotionClassifier(weights)predictions = classifier.classify(source)for filename, emotion in predictions.items():print(f'{filename}: {emotion}')