import gradio as gr
import tensorflow as tf

def classify_image(Image):
    model = tf.keras.models.load_model('tumor_model.h5')
    labels = [
        'Tumor - Glioma', 
        'No Tumor', 
        'Tumor - Meningioma', 
        'Tumor - Pituitary'
    ]
    
    inp = Image.reshape((-1, 128, 128, 3))
    inp = tf.keras.applications.resnet50.preprocess_input(inp)
    prediction = model.predict(inp).flatten()
    confidences = {labels[i]: float(prediction[i]) for i in range(4)}
    
    return confidences

image_input = gr.inputs.Image(shape = (128, 128))
label_output = gr.outputs.Label(num_top_classes = 4)

title = "Tumor Classification"
description = "Upload an image and get predictions for tumor classification."
examples = [
    ['Example 1.jpg'],
    ['Example 2.jpg']
]

interface = gr.Interface(
    fn = classify_image,
    inputs = image_input,
    outputs = label_output,
    title = title,
    description = description,
    examples = examples,
    button_style ='danger',
    theme = 'huggingface'
)

interface.launch()