Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import cv2 | |
| import torch | |
| from .exp_recognition_model import * | |
| from PIL import Image | |
| import os | |
| ############################################################################################################################# | |
| # Caution: Don't change any of the filenames, function names and definitions # | |
| ############################################################################################################################# | |
| current_path = os.path.dirname(os.path.abspath(__file__)) | |
| # ============================================================ | |
| # DEVICE | |
| # ============================================================ | |
| if torch.backends.mps.is_available(): | |
| device = torch.device("mps") | |
| elif torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| else: | |
| device = torch.device("cpu") | |
| # ============================================================ | |
| # FACE DETECTION | |
| # ============================================================ | |
| def detected_face(image): | |
| eye_haar = current_path + '/haarcascade_eye.xml' | |
| face_haar = current_path + '/haarcascade_frontalface_default.xml' | |
| face_cascade = cv2.CascadeClassifier(face_haar) | |
| eye_cascade = cv2.CascadeClassifier(eye_haar) | |
| gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) | |
| faces = face_cascade.detectMultiScale( | |
| gray, | |
| scaleFactor=1.3, | |
| minNeighbors=5 | |
| ) | |
| face_areas = [] | |
| images = [] | |
| required_image = 0 | |
| for i, (x, y, w, h) in enumerate(faces): | |
| face_cropped = gray[y:y+h, x:x+w] | |
| face_areas.append(w * h) | |
| images.append(face_cropped) | |
| required_image = images[np.argmax(face_areas)] | |
| required_image = Image.fromarray(required_image) | |
| return required_image | |
| # ============================================================ | |
| # LOAD MODEL | |
| # ============================================================ | |
| _model = None | |
| def load_model(): | |
| global _model | |
| if _model is None: | |
| model_path = current_path + '/expression_model.pth' | |
| checkpoint = torch.load( | |
| model_path, | |
| map_location=device | |
| ) | |
| model = facExpRec( | |
| num_classes=len(classes) | |
| ).to(device) | |
| if isinstance(checkpoint, dict): | |
| if 'model_state_dict' in checkpoint: | |
| model.load_state_dict( | |
| checkpoint['model_state_dict'] | |
| ) | |
| elif 'net_dict' in checkpoint: | |
| model.load_state_dict( | |
| checkpoint['net_dict'] | |
| ) | |
| else: | |
| model.load_state_dict(checkpoint) | |
| else: | |
| model.load_state_dict(checkpoint) | |
| model.eval() | |
| _model = model | |
| return _model | |
| # ============================================================ | |
| # EXPRESSION PREDICTION | |
| # ============================================================ | |
| def get_expression(img): | |
| model = load_model() | |
| face = detected_face(img) | |
| if face == 0: | |
| gray = cv2.cvtColor( | |
| img, | |
| cv2.COLOR_BGR2GRAY | |
| ) | |
| face = Image.fromarray(gray) | |
| image_tensor = trnscm(face) | |
| image_tensor = image_tensor.unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| outputs = model(image_tensor) | |
| probabilities = torch.softmax( | |
| outputs, | |
| dim=1 | |
| ) | |
| confidence, predicted = torch.max( | |
| probabilities, | |
| 1 | |
| ) | |
| predicted_idx = predicted.item() | |
| expression = classes[predicted_idx] | |
| return expression | |