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

teacher_model = from_pretrained_keras("Blazer007/consistency_training_with_supervision_teacher_model")

student_model = from_pretrained_keras("Blazer007/consistency_training_with_supervision_student_model")

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

IMG_SIZE = 72

def infer(input_image):
    print('#$$$$$$$$$$$$$$$$$$$$$$$$$ IN INFER $$$$$$$$$$$$$$$$$$$$$$$')
    # image_tensor = read_image(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))
    print(image_tensor.shape)
    predictions = teacher_model.predict(np.expand_dims((image_tensor), axis=0))
    print(predictions)
    predictions = np.squeeze(predictions)
    print(predictions)
    predictions = np.argmax(predictions) # , axis=2
    print(predictions)
    predicted_label = class_names[predictions.item()]
    print(predictions.item())
    print(predicted_label)
    return str(predicted_label)
    
input = gr.inputs.Image(shape=(IMG_SIZE, IMG_SIZE))
output = [gr.outputs.Label()]
examples = [[], []] 
title = "Image Classification using "
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>Space author: Vivek Rai</b> <br><b> Keras example author: Sayak Paul </b></p>"

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