import cv2
import gradio as gr
from predict_image import load_model, predict

def predict_fn(image, model_name):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    cv2.imwrite('./myimage.jpg', image)
    # model for emotion classification
    if model_name == 'EfficientNetB0':
        model_name = 'effb0'
    elif model_name == 'ResNet18':
        model_name = 'res18'
    else:
        raise ValueError('Enter correct model_name')
    
    model = load_model(model_name)
    
    out = predict('./myimage.jpg', './result.jpg', model)
    out = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
    return out


demo = gr.Interface(
    fn=predict_fn,
    inputs=[
        gr.inputs.Image(label="Input Image"),
        gr.Radio(['EfficientNetB0', 'ResNet18'], value='EfficientNetB0', label='Model Name')
        ],
    outputs=[
        gr.inputs.Image(label="Prediction"),
    ],
    title="Emotion Recognition Demo",
    description="Emotion Classification Model trained on FER Dataset",
    examples=[
        ["example/fear.jpg", 'EfficientNetB0'],
        ["example/sad.jpg", 'EfficientNetB0'],
        ["example/happy.jpg", 'EfficientNetB0'],
    ],
)

demo.launch(debug=True)