from transformers import ViTImageProcessor, AutoModelForImageClassification

import torch
import gradio as gr
import os
import glob
import mediapipe as mp
import numpy as np
from PIL import Image

feature_extractor = ViTImageProcessor.from_pretrained('ArdyL/VIT_SIBI_ALL')
model = AutoModelForImageClassification.from_pretrained('ArdyL/VIT_SIBI_ALL')


mp_drawing_styles = mp.solutions.drawing_styles
mp_holistic = mp.solutions.holistic
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils


def preprocess(im):
    with mp_holistic.Holistic(
            static_image_mode=True, min_detection_confidence=0.3, model_complexity=2) as hands:

        # Read image file with cv2 and process with face_mesh
        results = hands.process(im)
        image2 = np.array(im)
        annotated_image = image2.copy()
        annotated_image = np.empty(annotated_image.shape)
        annotated_image.fill(255)
        if results.pose_landmarks:
            mp_drawing.draw_landmarks(annotated_image, results.right_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
                                      mp_drawing.DrawingSpec(
                                          color=(0, 0, 0), thickness=2, circle_radius=2),
                                      mp_drawing.DrawingSpec(
                                          color=(0, 0, 0), thickness=2, circle_radius=2),
                                      )
            mp_drawing.draw_landmarks(annotated_image, results.left_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
                                      mp_drawing.DrawingSpec(
                                          color=(0, 0, 0), thickness=2, circle_radius=2),
                                      mp_drawing.DrawingSpec(
                                          color=(0, 0, 0), thickness=2, circle_radius=2),
                                      )
            mp_drawing.draw_landmarks(annotated_image, results.face_landmarks, mp_holistic.FACEMESH_TESSELATION,
                                      landmark_drawing_spec=None,
                                      connection_drawing_spec=mp_drawing_styles
                                      .get_default_face_mesh_tesselation_style()
                                      )
            mp_drawing.draw_landmarks(annotated_image,  results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
                                      mp_drawing.DrawingSpec(
                                          color=(0, 0, 0), thickness=2, circle_radius=2),
                                      mp_drawing.DrawingSpec(
                                          color=(0, 0, 0), thickness=2, circle_radius=2),
                                      )

            annotated_image[...] /= 255
    return annotated_image


def classify_image(image):
    preprocessedImage = preprocess(image)
    with torch.no_grad():
        model.eval()

        inputs = feature_extractor(
            images=preprocessedImage, return_tensors="pt")
        outputs = model(**inputs)

    logits = outputs.logits
    predicted_label = logits.argmax(-1).item()
    label = model.config.id2label[predicted_label]
    # prob = torch.nn.functional.softmax(logits, dim=1)

    # top10_prob, top10_indices = torch.topk(prob, 5)
    # top10_confidences = {}
    # for i in range(5):
    #     top10_confidences[model.config.id2label[int(
    #         top10_indices[0][i])]] = float(top10_prob[0][i])

    return label  # confidences


with gr.Blocks(title=">ViT - SIBI Classifier"
               ) as demo:
    # gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">ViT - SIBI Classifier</div>""")

    # txt_2 = gr.Textbox(label="Input 2", visible=False)
    with gr.Tab("Upload Image", id='upload-image', ):
        with gr.Row():
            uploadImage = gr.Image(
                type="numpy", image_mode="RGB", shape=(224, 224))
            output_label = gr.Label(label="Hasil", num_top_classes=5)
        with gr.Row():
            send_btn = gr.Button("Terjemahkan")
            send_btn.click(fn=classify_image, inputs=uploadImage,
                           outputs=output_label)
        # with gr.Row():
        #     with gr.Column(scale=1):
        #         gr.Examples([
        #             # Alie
        #             ['./samples/baik.jpg', 'Baik'],
        #             # ['./samples/dengar.jpg', 'Dengar'],
        #             # ['./samples/diam.jpg', 'Diam'],
        #             # ['./samples/ibu.jpg', 'Ibu'],
        #             # Hardy
        #             ['./samples/maaf.jpg', 'Maaf'],
        #             # ['./samples/mahasiswa.jpg', 'Mahasiswa'],
        #             # ['./samples/pilih.jpg', 'Pilih'],
        #             # ['./samples/punya.jpg', 'Punya'],
        #             # ['./samples/saya.jpg', 'Saya'],
        #             # ['./samples/sehat.jpg', 'Sehat'],
        #             # ['./samples/zuhur.jpg', 'Zuhur'],

        #         ],
        #             [uploadImage, txt_2])
        #   with gr.Column(scale=1):
        #     gr.Label('Dengar')
        # with gr.Row():
        #   with gr.Column(scale=1):
        #     gr.Examples(['./baik.jpg'], inputs=uploadImage)

    with gr.Tab("Capture Image", id='capture-image'):

        with gr.Row():
            streamImage = gr.Image(
                type="numpy", source='webcam', image_mode="RGB", shape=(224, 224))
            output_label2 = gr.Label(label="Hasil", num_top_classes=5)
        with gr.Row():
            send_btn2 = gr.Button("Terjemahkan")
            send_btn2.click(fn=classify_image,
                            inputs=streamImage, outputs=output_label2)

            # with gr.Row():
            #     gr.Examples(['./dengar.jpg'], label='Sample images : cat', inputs=input_image)
            #     gr.Examples(['./baik.jpg'], label='cheetah', inputs=input_image)
            # gr.Examples(['./samples/hotdog.jpg'], label='hotdog', inputs=input_image)
            # gr.Examples(['./samples/lion.jpg'], label='lion', inputs=input_image)
            # gr.Examples(example_files, inputs=input_image)

# demo.queue(concurrency_count=3)
demo.launch(debug=True)