'''from pathlib import Path
import shutil
import itertools
import os, cv2, numpy as np'''
import gradio as gr
import torch
from transformers import AutoModelForImageClassification
from optimum.pipelines import pipeline
from PIL import Image
import numpy as np
device = 1 if torch.cuda.is_available() else "cpu"

# chk_point = "kdhht2334/autotrain-diffusion-emotion-facial-expression-recognition-40429105176"

model = AutoModelForImageClassification.from_pretrained("./autotrain-diffusion-emotion-facial-expression-recognition-40429105176")

##Add face detector
from facenet_pytorch import MTCNN, InceptionResnetV1
mtcnn = MTCNN(image_size=300, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True)
resnet = InceptionResnetV1(pretrained='vggface2').eval()

emotion_dict = {
'neutral': '0',
'happy': '1',
'sad' :'2',
'surprise': '3',
'fear': '4',
'disgust': '5',
'angry': '6',
'uncertain': '7',
'nonface': '8',
}


output_img_size = (2100, 700)


try:
    pipe = pipeline(
        "image-classification",
        model,
        accelerator="bettertransformer",
        device=device,
    )
except NotImplementedError:
    from transformers import pipeline

    pipe = pipeline("image-classification", model, device=device)


def face_detector(input_img):
    img = Image.fromarray(input_img)
    bbox, _ = mtcnn.detect(img)
    bbox = bbox.squeeze().tolist()
    crop = img.crop(bbox)
    return crop

def predict(image):
    cropped_face = face_detector(image)
    face_w, face_h = cropped_face.size
    face_re_w = int(face_w * (700 / face_h))
    resized_face = cropped_face.resize((face_re_w, 700))
    output_img = Image.new("RGBA", output_img_size)
    output_img.paste(resized_face, (1050 - int(face_re_w/2), 0))
    predictions = pipe(cropped_face)
    return output_img, {p["label"]: p["score"] for p in predictions}

gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Upload image"),
    outputs=["image", "label"],
    examples=[["examples/happy.png"], ["examples/angry.png"], ["examples/surprise.png"]],
    title="Demo - DiffusionFER",
).launch()