Spaces:
Sleeping
Sleeping
File size: 3,188 Bytes
959739d |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
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
|