import gradio as gr
import tensorflow as tf
from keras_tuner import HyperParameters
from src.models import MakeHyperModel
from src.preprocessing import get_data_augmentation

from src.config import IMAGE_SIZE

data_augmentation = get_data_augmentation()

latest = tf.train.latest_checkpoint('./tuner_model/cat-vs-dog/trial_0484d8d758a5ef7b91ca97d334ba7870/checkpoints/epoch_0')
hypermodel = MakeHyperModel(input_shape=IMAGE_SIZE + (3,), num_classes=2, data_augmentation=data_augmentation)
model = hypermodel.build(hp=HyperParameters())
model.load_weights(latest).expect_partial()

def cat_vs_dog(image):
    img_array = tf.constant(image, dtype=tf.float32)
    img_array = tf.expand_dims(img_array, 0)
    predictions = model.predict(img_array)
    score = predictions[0]
    return {'cat': float((1 - score)), 'dog': float(score)}

iface = gr.Interface(
    cat_vs_dog, 
    gr.inputs.Image(shape=IMAGE_SIZE), 
    gr.outputs.Label(num_top_classes=2),
    capture_session=True,
    interpretation="default",
    examples=[
        ["examples/cat1.jpg"],
        ["examples/cat2.jpg"],
        ["examples/dog1.jpeg"],
        ["examples/dog2.jpeg"]
    ])

if __name__ == "__main__":
    iface.launch()