Spaces:
Sleeping
Sleeping
import cv2 | |
from utils import read_image,get_valid_augs | |
import torch | |
import torch.nn.functional as F | |
from model import ResNet50Model | |
import numpy as np | |
CKPT_EMOTION = 'emotion_resnet_model' | |
CKPT_CELEBRITY = 'celebrity_resnet_model' | |
FaceInverseTargetMapper = {0: 'Tom Hanks',1: 'Sandra Bullock',2: 'Natalie Portman',3: 'Scarlett Johansson',4: 'Robert Downey Jr',5: 'Nicole Kidman',6: 'Brad Pitt', | |
7: 'Hugh Jackman',8: 'Tom Cruise',9: 'Leonardo DiCaprio',10: 'Megan Fox',11: 'Johnny Depp',12: 'Will Smith',13: 'Denzel Washington',14: 'Jennifer Lawrence',15: 'Kate Winslet',16: 'Angelina Jolie'} | |
EmotionMapper = {0: 'sadness',1: 'contempt',2: 'happiness',3: 'surprise', | |
4: 'fear',5: 'anger',6: 'disgust',7: 'neutrality'} | |
def predict_one_image(path) : | |
image = read_image(path) | |
image = get_valid_augs()(image=image)['image'] | |
image = torch.tensor(image,dtype=torch.float) | |
image = image.reshape((1,3,224,224)) | |
emotion_model = ResNet50Model('emotion') | |
#loading ckpt | |
emotion_model.load_state_dict(torch.load(CKPT_EMOTION,map_location=torch.device('cpu'))) | |
celebrity_model = ResNet50Model('celebrity') | |
#loading ckpt | |
celebrity_model.load_state_dict(torch.load(CKPT_CELEBRITY,map_location=torch.device('cpu'))) | |
with torch.no_grad() : | |
#emotion | |
outputs = emotion_model(image) | |
outputs = torch.nn.functional.softmax(outputs).cpu().detach().numpy() | |
print(outputs.shape) | |
emotion = np.argmax(outputs,axis=1)[0] | |
emotion_proba = np.max(outputs,axis=1)[0] | |
print(emotion_proba) | |
#celebrity | |
outputs = celebrity_model(image) | |
outputs = torch.nn.functional.softmax(outputs).cpu().detach().numpy() | |
print(outputs.shape) | |
celebrity = np.argmax(outputs,axis=1)[0] | |
celebrity_proba = np.max(outputs,axis=1)[0] | |
print(celebrity_proba) | |
if celebrity_proba<0.45 : | |
return f"Unkonwn Person Detected with emotion {EmotionMapper[emotion]} " | |
return f"Detected {FaceInverseTargetMapper[celebrity]} with emotion {EmotionMapper[emotion]} " |