import os
import tensorflow as tf
from tensorflow import keras
import gradio as gr

# Suppress TensorFlow logs
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# Load the model
model_path = 'doctor_ai_model.h5'  # Change to your model file
model = keras.models.load_model(model_path)

def predict(input_data):
    # Convert input to tensor
    input_tensor = tf.convert_to_tensor(input_data)

    # Ensure the input shape is correct
    if input_tensor.shape[1] != 27:  # Adjust based on your input shape
        return "Input data must have shape: (None, 27)"
    
    # Expand dimensions for the model
    input_tensor = tf.expand_dims(input_tensor, axis=0)

    # Make prediction
    prediction = model.predict(input_tensor)
    predicted_class = tf.argmax(prediction, axis=1).numpy().tolist()
    
    return predicted_class

# Create Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Dataframe(type='numpy', label='Input Data (should have shape (None, 27))'),
    outputs='json'
)

# Launch the Gradio app
iface.launch()