import gradio as gr
from tensorflow.keras.models import load_model
from PIL import Image, ImageOps
import numpy as np
from tensorflow.keras.layers import DepthwiseConv2D

def custom_depthwise_conv2d(*args, **kwargs):
    if 'groups' in kwargs:
        del kwargs['groups']  # Retirer 'groups'
    return DepthwiseConv2D(*args, **kwargs)

# Load model
model = load_model("keras_model.h5", custom_objects={'DepthwiseConv2D': custom_depthwise_conv2d}, compile=False)

# Load labels
with open("labels.txt", "r") as file:
    class_names = file.readlines()

# Create predict fonction
def predict(image):
    image = ImageOps.fit(image, (224, 224), Image.Resampling.LANCZOS)
    
    image_array = np.asarray(image)
    
    normalized_image_array = (image_array.astype(np.float32) / 127.5) - 1
    
    data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
    data[0] = normalized_image_array

    # Make prediction
    prediction = model.predict(data)
    index = np.argmax(prediction)
    class_name = class_names[index].strip() 
    confidence_score = prediction[0][index]
    
    return class_name, confidence_score

# Créer l'interface Gradio
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Download image"),
    outputs=[
        gr.Label(label="Class predict"),
        gr.Number(label="Confidence Score")
    ],
    title="Medical Assistant",
    description="Upload a dental image and our app will predict its class."
)

iface.launch()