from PIL import Image
import tensorflow as tf
import numpy as np
import gradio as gr
import io
import json

# Load the model
model_path = 'final_teath_classifier.h5'
model = tf.keras.models.load_model(model_path)

# Define preprocessing function


# Define prediction function
def predict_image(image):
    # Save the image to a file-like object
    image_bytes = io.BytesIO()
    image.save(image_bytes, format="JPEG")

    # Load the image from the file-like object
    image = tf.keras.preprocessing.image.load_img(image_bytes, target_size=(256, 256,3))
    image = np.array(image)/255
    image = np.expand_dims(image, axis=0)

    # Make a prediction
    prediction = model.predict(image)

    # Get the probability of being 'Clean' or 'Carries'
    probabilities = tf.nn.softmax(prediction, axis=-1)
    predicted_class_index = np.argmax(probabilities)
    if predicted_class_index == 1:
        predicted_label = "Clean"
        predicted_probability = probabilities[0][1] * 100  # Convert to percentage
    elif predicted_class_index == 0:
        predicted_label = "Carries"
        predicted_probability = probabilities[0][0] * 100  # Convert to percentage

    # Return the prediction result as a dictionary
    return {"Predicted Label": predicted_label}


# Create the interface
input_interface = gr.Image(type="pil")
output_interface = "json"

iface = gr.Interface(
    fn=predict_image,
    inputs=input_interface,
    outputs=output_interface)

# Launch the interface
iface.launch(share=True)