import gradio as gr
import glob
import os

from src.core import load_model, predict_traits

TRAIT_NAMES = [
    "Extraversion",
    "Agreeableness",
    "Conscientiousness",
    "Neurotisicm",
    "Openness",
]


def get_traits(video):
    model = load_model()
    # if webcam_video:
    #     trait_values = predict_traits(webcam_video, model)
    # else:
    trait_values = predict_traits(video, model)
    return {k: float(v) for k, v in zip(TRAIT_NAMES, trait_values)}


demo = gr.Interface(
    get_traits,
    inputs=gr.Video(label="Video", include_audio=True),
    outputs=gr.Label(num_top_classes=5, label="Results"),
    title="Personality Traits Prediction [Prototype]",
    description="Predicts the 5 psychological traits of a person using an short introduction video",
    thumbnail="https://cdn-icons-png.flaticon.com/512/3392/3392044.png",
    examples="egs",
    cache_examples=True,
)
demo.launch()