import gradio as gr
import PIL 
from PIL import Image
import tensorflow
import numpy
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model

model = load_model('vgg_model.h5')

def preprocess_image(img):
    
    img = Image.fromarray(img)

    # Resize the image to match the model's expected input size
    img = img.resize((150, 150))

    # Convert the image to a numpy array
    img_array = image.img_to_array(img)

    # Expand the dimensions to match the model's expected input shape
    img_array = numpy.expand_dims(img_array, axis=0)

    # Normalize the pixel values to be in the range [0, 1]
    img_array /= 255.0

    return img_array

def predict(img):
    # Preprocess the image
    img_array = preprocess_image(img)

    # Make predictions
    prediction = model.predict(img_array)

    prediction = round(float(prediction[0][0]), 5)
    # Create a dictionary of confidences
    # confidences = {"Cat": prediction, 'Dog': 1 - prediction}

    # Return a tuple with the predicted label and confidence
    if prediction >= 0.47:
        # return {'Dog': prediction, "Cat": 1-prediction}
        return {'Dog': 1}
        
    else:
        # return {'Dog': prediction, "Cat": 1-prediction}
        return {"Cat": 1}
        

    # return confidences

title = "Cats🐱 - Dogs🐶 Classification 🐾" 
description = "A VGG model to classify images of pets as cats or dogs."

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=1),
    description=description,
    title= title
)


demo.launch(share=True)