from random import choices
import numpy as np
import gradio as gr
from glob import glob
import tensorflow as tf
from tensorflow import keras

# Model & Pre-requisites
model_path = './FastFood.keras'
ffc = keras.models.load_model(model_path, compile=False)

class_names_path = './Fast Food-ClassNames.txt'
class_names = []
with open(class_names_path, mode='r') as f:
    class_names = f.read().split(',')[:-1]

# Utility Functions
def predict_fast_food(image, labels=class_names, model=ffc):
    image = tf.cast(image, tf.float32)

    if image.shape[-2]!=224:
        image = tf.image.resize(image, (224,224))

    if np.max(image)==255:
        image = image/255.

    if len(image.shape) == 3:
        image = tf.squeeze(image)[tf.newaxis, ...]

        pred_proba = model.predict(image, verbose=0)[0]
        label = tf.argmax(pred_proba, axis=-1)
        pred_class = labels[int(label)]
        return pred_class, pred_proba[label]
    else:
        pred_probas = model.predict(image, verbose=0)
        labels = tf.argmax(pred_probas, axis=-1)
        pred_classes = [class_names[label] for label in labels]
        probas = tf.math.reduce_max(pred_probas, axis=-1)
        return pred_classes, probas

def load_image(image_path):
    image = tf.io.read_file(image_path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, (224,224))
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = image/255.
    return image

# Load Example Images
subset_ds_path = './Fast FoodSubset'

# Select 5 images per class
example_image_paths = []
for class_ss_path in glob(subset_ds_path + '/*'):
    image_paths = glob(class_ss_path + '/*')
    selected_images = choices(image_paths, k=5)
    example_image_paths.extend(selected_images)

example_images = [load_image(path).numpy() for path in example_image_paths]

# Define Interface
with gr.Blocks(theme='ocean') as app:
    
    # Title or header (optional)
    gr.Markdown("### 🍔 Fast Food Classifier Demo")
    
    # Take Image Input
    image_input = gr.Image(label='Image Input')

    # Prediction Button 
    pred_btn = gr.Button('Predict')
    
    # 2 Outputs
    with gr.Row():

        # Output of the Predicted Class
        class_out = gr.Textbox(label='Predicted Class', placeholder='Hmm... Looking for something yummy.')
        proba_out = gr.Textbox(label='Predicted Class Probability', placeholder='I believe on myself but numbers don\'t lie.')

    # Add example images
    gr.Examples(
        examples=example_images,
        inputs=image_input,
        label="Try these example images"
    )

    def predict_fast_food_wrapper(image):
        class_label, proba = predict_fast_food(image)
        return class_label, f'{proba:.3%}'
    
    # On Click Action
    pred_btn.click(
        fn=predict_fast_food_wrapper,
        inputs=image_input,
        outputs=[class_out, proba_out]
    )
    
if __name__ == '__main__':
    # Launch Application
    app.launch()