import numpy as np
import tensorflow as tf
import gradio as gr
from huggingface_hub import from_pretrained_keras

teacher_model = from_pretrained_keras("keras-io/consistency_training_with_supervision_teacher_model")

student_model = from_pretrained_keras("keras-io/consistency_training_with_supervision_student_model")

class_names = [
    "Airplane",
    "Automobile",
    "Bird",
    "Cat",
    "Deer",
    "Dog",
    "Frog",
    "Horse",
    "Ship",
    "Truck",
]

examples = [
    ['./aeroplane.png'],
    ['./horse.png'],
    ['./ship.png'],
    ['./truck.png']
] 

IMG_SIZE = 72

def teacher_model_output(image_tensor):
    predictions = teacher_model.predict(np.expand_dims((image_tensor), axis=0))
    predictions = np.squeeze(predictions)
    predictions = np.argmax(predictions)
    predicted_label = class_names[predictions.item()]
    return str(predicted_label)
    
def student_model_output(image_tensor):
    predictions = student_model.predict(np.expand_dims((image_tensor), axis=0))
    predictions = np.squeeze(predictions)
    predictions = np.argmax(predictions)
    predicted_label = class_names[predictions.item()]
    return str(predicted_label)

def infer(input_image):
    image_tensor = tf.convert_to_tensor(input_image)
    image_tensor.set_shape([None, None, 3])
    image_tensor = tf.image.resize(image_tensor, (IMG_SIZE, IMG_SIZE))
    return teacher_model_output(image_tensor), student_model_output(image_tensor)
    
input = gr.inputs.Image(shape=(IMG_SIZE, IMG_SIZE))
output = [gr.outputs.Label(label = "Teacher Model Output"), gr.outputs.Label(label = "Student Model Output")]

title = "Image Classification using Consistency training with supervision"
description = "Upload an image or select from examples to classify it.<br>The allowed classes are - Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, Horse, Ship, Truck.<br><p><b>Teacher Model Repo - https://huggingface.co/keras-io/consistency_training_with_supervision_teacher_model</b> <br><b> Student Model Repo - https://huggingface.co/keras-io/consistency_training_with_supervision_student_model </b><br><b>Keras Example - https://keras.io/examples/vision/consistency_training/</b></p>"


article = "<div style='text-align: center;'><a href='https://twitter.com/_Blazer_007' target='_blank'>Space by Vivek Rai</a><br><a href='https://twitter.com/RisingSayak' target='_blank'>Keras example by Sayak Paul</a></div>"

gr_interface = gr.Interface(
    infer, 
    input, 
    output, 
    examples=examples, 
    allow_flagging=False, 
    analytics_enabled=False, 
    title=title, 
    description=description,
    article=article).launch(enable_queue=True, debug=True)