app/app_utils.py

"""
File: app_utils.py
Author: Elena Ryumina and Dmitry Ryumin
Description: This module contains utility functions for facial expression recognition application.
License: MIT License
"""

import torch
import numpy as np
import mediapipe as mp
from PIL import Image

# Importing necessary components for the Gradio app
from app.model import pth_model, pth_processing
from app.face_utils import get_box
from app.config import DICT_EMO


mp_face_mesh = mp.solutions.face_mesh


def preprocess_and_predict(inp):
    inp = np.array(inp)

    if inp is None:
        return None, None

    try:
        h, w = inp.shape[:2]
    except Exception:
        return None, None

    with mp_face_mesh.FaceMesh(
        max_num_faces=1,
        refine_landmarks=False,
        min_detection_confidence=0.5,
        min_tracking_confidence=0.5,
    ) as face_mesh:
        results = face_mesh.process(inp)
        if results.multi_face_landmarks:
            for fl in results.multi_face_landmarks:
                startX, startY, endX, endY = get_box(fl, w, h)
                cur_face = inp[startY:endY, startX:endX]
                cur_face_n = pth_processing(Image.fromarray(cur_face))
                prediction = (
                    torch.nn.functional.softmax(pth_model(cur_face_n), dim=1)
                    .detach()
                    .numpy()[0]
                )
                confidences = {DICT_EMO[i]: float(prediction[i]) for i in range(7)}

    return cur_face, confidences