import tensorflow as tf
import numpy as np
import gradio as gr
from tensorflow.keras.preprocessing import image

# Load the trained model
model = tf.keras.models.load_model("ecg_classification_model (1).keras", compile=False)

# Class labels (modify based on your dataset)
class_labels = [
    "Left Bundle Branch Block",
    "Normal",
    "Premature Atrial Contraction",
    "Premature Ventricular Contractions",
    "Right Bundle Branch Block",
    "Ventricular Fibrillation"
]

# Function to preprocess the image
def preprocess_image(img):
    img = img.resize((224, 224))  # Resize to match model input
    img_array = np.array(img) / 255.0  # Normalize
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    return img_array

# Function to make a prediction
def predict_ecg(img):
    processed_img = preprocess_image(img)
    prediction = model.predict(processed_img)
    predicted_class = class_labels[np.argmax(prediction)]
    return f"Predicted Class: {predicted_class}"

# Create Gradio Interface
iface = gr.Interface(
    fn=predict_ecg,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="ECG Image Classifier",
    description="Upload an ECG image to classify it."
)

# Run the app
if __name__ == "__main__":
    iface.launch(share=True)