import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import gradio as gr
import numpy as np
import cv2
import tensorflow as tf

# Load the model
model = tf.keras.models.load_model('model.h5')

def preprocess_image(image):
    resized_img = cv2.resize(image, (180, 180))
    img_array = np.array(resized_img).reshape((1, 180, 180, 3))
    return img_array

def predict_pneumonia(image):
    img_array = preprocess_image(image)
    prediction = model.predict(img_array)[0][0]

    # Use a more robust threshold for determining whether an image has pneumonia
    threshold = 0.5

    if prediction >= threshold:
        return "Pneumonia"
    else:
        return "Normal"

inputs = gr.inputs.Image(shape=(180, 180))
outputs = gr.outputs.Label(num_top_classes=2)

gradio_interface = gr.Interface(
    fn=predict_pneumonia,
    inputs=inputs,
    outputs=outputs,
    title="Pneumonia X-Ray Classification API",
    description="This API classifies images of chest X-rays as having pneumonia or being normal.",
    examples=[
        ["person1000_bacteria_2931.jpeg"],
        ["person1000_virus_1681.jpeg"],
        ["person1946_bacteria_4875.jpeg"],
        ["person1952_bacteria_4883.jpeg"],
        ["NORMAL2-IM-1427-0001.jpeg"]
    ],
    theme="default",
    allow_flagging=False
)

gradio_interface.launch()