Spaces:
Sleeping
Sleeping
| from torchvision import models | |
| import torch.nn as nn | |
| import torch | |
| import torch.nn.functional as F | |
| import torchvision.transforms as transforms | |
| from PIL import Image | |
| import cv2 | |
| def create_emotion_model(num_ftrs, num_emotions): | |
| return nn.Sequential( | |
| nn.Linear(num_ftrs + num_emotions, 128), | |
| nn.ReLU(), | |
| nn.Linear(128, 64), | |
| nn.ReLU(), | |
| nn.Linear(64, 2), | |
| ) | |
| def load_models(val_model_path,val_featmodel_path): | |
| 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]), | |
| ] | |
| ) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| resnet = models.resnet18(pretrained=False) | |
| num_ftrs = resnet.fc.in_features | |
| resnet.fc = nn.Identity() | |
| resnet.load_state_dict( | |
| torch.load( | |
| val_featmodel_path, | |
| map_location=device | |
| ) | |
| ) | |
| resnet = resnet.to(device) | |
| # num_ftrs = resnet.fc.in_features | |
| num_emotions = 1 | |
| emotion_model = create_emotion_model(num_ftrs, num_emotions).to(device) | |
| emotion_model.load_state_dict( | |
| torch.load( | |
| val_model_path, | |
| map_location=device | |
| ) | |
| ) | |
| return resnet,emotion_model | |
| def va_predict(emotion_model,resnet,faces,emotions): | |
| 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]), | |
| ] | |
| ) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def model_forward(images, emotions): | |
| resnet_features = resnet(images) | |
| batch_size = resnet_features.size(0) | |
| emotions = emotions.view(batch_size, -1) | |
| x = torch.cat((resnet_features, emotions), dim=1) | |
| output = emotion_model(x) | |
| return output | |
| arousal_list = [] | |
| valence_list = [] | |
| stress_list = [] | |
| from tqdm import tqdm | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| for face, emotion in tqdm(zip(faces, emotions)): | |
| if face is not None: | |
| face_pil = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB)) | |
| face_tensor = transform(face_pil).unsqueeze(0).to(device) | |
| # print(emotion) | |
| # print(emotion) | |
| emotion = emotion.to(device) | |
| output_va = model_forward(face_tensor, emotion) | |
| arousal = output_va[0][0].item() | |
| norm_arousal = float(output_va[0][0].item()) / 2 + 0.5 | |
| valence = output_va[0][1].item() | |
| norm_valence = float(output_va[0][1].item()) / 2 + 0.5 | |
| stress = (1 - norm_valence) * norm_arousal | |
| arousal_list.append(arousal) | |
| valence_list.append(valence) | |
| stress_list.append(stress) | |
| else: | |
| arousal_list.append('frame error') | |
| valence_list.append('frame error') | |
| stress_list.append('frame error') | |
| return valence_list, arousal_list, stress_list | |