import os
import gradio as gr
import cv2
from insightface.app import FaceAnalysis
from hsemotion_onnx.facial_emotions import HSEmotionRecognizer


def resize(image, target_size):
    # Get the dimensions of the input image
    height, width = image.shape[0], image.shape[1]

    # Calculate the scaling factor needed to resize the image to the target size
    scaling_factor = min(target_size[0] / width, target_size[1] / height)
    # Resize the image using cv2.resize
    resized_image = cv2.resize(image, None, fx=scaling_factor, fy=scaling_factor, interpolation=cv2.INTER_LINEAR)

    return resized_image


def facial_emotion_recognition(img):

    img = resize(img, target_size)
    
    faces = face_detector.get(img)

    if len(faces) > 0:

        highest_score_box = (0, 0, 0, 0)  # x, y, w, h
        highest_score = 0

        for face in faces:
            if face['det_score'] > highest_score:
                highest_score = face['det_score']
                x1, y1, x2, y2 = face['bbox'].astype(int)
                x_margin = int((x2 - x1) * face_margin)
                y_margin = int((y2 - y1) * face_margin)
                x = max(0, x1 - x_margin)
                y = max(0, y1 - y_margin)
                w = min(x2 + x_margin, img.shape[1]) - x
                h = min(y2 + y_margin, img.shape[0]) - y
                highest_score_box = (x, y, w, h)

        x, y, w, h = highest_score_box
        emotion, _ = hse_emo_model.predict_emotions(img[y:y+h, x:x+w], logits=True)

        cv2.rectangle(img, (x, y), (x+w, y+h), (0, 0, 255), 2)
        cv2.putText(img, emotion, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)

    return img

face_margin = 0.1
target_size = (640, 640)
model_name = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'buffalo_sc')
face_detector = FaceAnalysis(name=model_name, allowed_modules=['detection'], providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
face_detector.prepare(ctx_id=0, det_size=(640, 640))


hse_emo_model = HSEmotionRecognizer(model_name='enet_b0_8_best_vgaf')

webcam = gr.Image(image_mode='RGB', type='numpy', source='webcam', label='Input Image')
webcam_output = gr.Image(image_mode='RGB', type='numpy', label='Output Image')
webcam_interface = gr.Interface(facial_emotion_recognition, inputs=webcam, outputs=webcam_output)

upload = gr.Image(image_mode='RGB', type='numpy', source='upload', label='Input Image')
upload_output = gr.Image(image_mode='RGB', type='numpy', label='Output Image')
upload_interface = gr.Interface(facial_emotion_recognition, inputs=upload, outputs=upload_output, examples='examples')

demo = gr.TabbedInterface(interface_list=[upload_interface, webcam_interface], tab_names=['Upload', 'Webcam'])
demo.launch()