import json
import numpy as np
import gradio as gr
import tensorflow as tf
from PIL import Image
from tensorflow.keras.models import load_model
# Load the model
model_path = 'final_teath_classifier.h5'
model = tf.keras.models.load_model(model_path)

# Define preprocessing function
def preprocess_image(image):
    # Resize the image to match input size
    image = Image.fromarray(image)
    image = image.resize((256, 256))
    # Convert image to array and preprocess input
    img_array = np.array(image) / 255.0
    # Add batch dimension
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Define prediction function
def predict_image(image):
    img_array = preprocess_image(image)
    outputs = model(img_array)
    predictions = tf.nn.softmax(outputs.logits, axis=-1)
    predicted_class = np.argmax(predictions)
    if predicted_class == 0:
        predict_label = "Clean"
    else:
        predict_label = "Carries"
    return {"prediction": predict_label, "confidence": float(np.max(predictions))}

# Create the interface
input_interface = gr.inputs.Image(shape=(256, 256), image_mode='RGB')
output_interface = gr.outputs.Label(num_top_classes=2)

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

# Launch the interface
iface.launch()