from keras.applications.vgg16 import VGG16,preprocess_input, decode_predictions
import keras.utils as image
from tensorflow.keras.layers import Dense, Flatten ,Dropout
from tensorflow.keras import Model
import numpy as np
import gradio as gr
import tensorflow as tf

vgg16 = VGG16(weights='imagenet', input_shape=(224,224,3), classes=10,include_top=False)

for layer in vgg16.layers:
            layer.trainable = False
            
x = Flatten()(vgg16.output)
x = Dense(256, activation='relu')(x)
x = Dropout(0.5)(x)
predictions = Dense(10, activation='softmax')(x)

vgg16_model = Model(inputs=vgg16.input, outputs=predictions)

vgg16_model.load_weights("VGG16_model.h5")

def prediction(input_image):
    # img = image.load_img(input_image, target_size=(224, 224))
    # x = image.img_to_array(img)
    img = tf.image.resize(input_image,(224,224))
    x = np.expand_dims(img, axis=0)
    # x = preprocess_input(x)
    preds=vgg16_model.predict(x)
    # create a list containing the class labels
    # find the index of the class with maximum score
    pred = np.argmax(preds, axis=-1)
    
    class_names = ['dog','horse','elephant','butterfly','chicken','cat','cow','sheep','spider','squirrel']
    # print the label of the class with maximum score
    return class_names[pred[0]]
    
# animals_classes = prediction("OIF-e2bexWrojgtQnAPPcUfOWQ.jpeg")

gr.Interface(fn=prediction,inputs=gr.Image(),outputs=gr.Label(num_top_classes=1)).launch()