from huggingface_hub import hf_hub_download
import tensorflow as tf
import gradio as gr
import numpy as np
from PIL import Image
import os

# Disable GPU usage
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

# Download and load model
model_path = hf_hub_download(repo_id="Owos/tb-classifier", filename="tb_model.h5")
model = tf.keras.models.load_model(model_path)

# Inference function
def predict_tb(img: Image.Image):
    try:
        image = img.convert("RGB").resize((224, 224))
        image_array = np.array(image) / 255.0
        image_array = image_array[np.newaxis, ...]
        prediction = model.predict(image_array)[0][0]
        label = "🦠 Tuberculosis Detected" if prediction > 0.5 else "🫁 Normal"
        confidence = prediction if prediction > 0.5 else 1 - prediction
        return f"{label} (Confidence: {confidence:.2%})"
    except Exception as e:
        return f"❌ Error during prediction: {str(e)}"

# Gradio UI
iface = gr.Interface(
    fn=predict_tb,
    inputs=gr.Image(type="pil", label="Upload Chest X-ray Image"),
    outputs="text",
    title="🩻 Tuberculosis Detection from Chest X-ray",
    description="Upload a chest X-ray to detect signs of Tuberculosis using an AI model (ResNet50). For educational & demo use only."
)

# Launch the app
iface.launch()