import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Make sure to decorate the custom metric class with @register_keras_serializable
@tf.keras.utils.register_keras_serializable()
class F1Score(tf.keras.metrics.Metric):
    def __init__(self, name='f1_score', **kwargs):
        super().__init__(name=name, **kwargs)
        self.precision = tf.keras.metrics.Precision()
        self.recall = tf.keras.metrics.Recall()

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.round(y_pred)
        self.precision.update_state(y_true, y_pred, sample_weight)
        self.recall.update_state(y_true, y_pred, sample_weight)

    def result(self):
        p = self.precision.result()
        r = self.recall.result()
        return 2 * ((p * r) / (p + r + tf.keras.backend.epsilon()))

    def reset_states(self):
        self.precision.reset_states()
        self.recall.reset_states()

# Load your TensorFlow model
model_path = './ACC0.9322_valACC0.9148_loss0.3592_Epoch83.keras'
model = tf.keras.models.load_model(model_path, custom_objects={'F1Score': F1Score})

def predict(image):
    # Preprocess the image to match the input shape of the model
    img = Image.fromarray(image.astype('uint8'), 'RGB')
    img = img.resize((224, 224))  # Update the size if different
    img_array = np.array(img)
    img_array = img_array / 255.0  # Normalize if your model expects normalization
    img_array = img_array[np.newaxis, ...]  # Model expects a batch dimension

    # Make prediction
    prediction = model.predict(img_array)
    pred_class = 'Pneumonia' if prediction[0][0] > 0.5 else 'Normal'  # Update based on your model output
    return {pred_class: float(prediction[0][0])}

iface = gr.Interface(fn=predict,
                     inputs=gr.components.Image(image_mode='RGB', label="Upload X-ray Image"),
                     outputs=gr.components.Label(num_top_classes=2),
                     title="Pneumonia Detection",
                     description="Upload an X-ray image to detect pneumonia.")

if __name__ == "__main__":
    iface.launch()