import gradio as gr
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image

# Load the model
print("Loading MobileNetV2 model...")
try:
    model = load_model("model.keras")
    print("✅ Model loaded successfully!")
except Exception as e:
    print(f"❌ Model loading failed: {e}")
    model = None

class_names = [
    "Oral Homogenous Leukoplakia",
    "Oral Non-Homogenous Leukoplakia",
    "Other Oral White Lesions"
]

def predict_image(img):
    if model is None:
        return "Error: Model failed to load. Please check if model.keras is uploaded.", {}

    try:
        # Preprocess image
        img = img.resize((224, 224))
        img_array = image.img_to_array(img)
        img_array = np.expand_dims(img_array, axis=0)
        img_array = img_array / 255.0

        # Predict
        predictions = model.predict(img_array, verbose=0)
        predicted_class = int(np.argmax(predictions[0]))
        confidence = float(np.max(predictions[0]) * 100)

        result = class_names[predicted_class]
        confidences = {class_names[i]: round(float(predictions[0][i] * 100), 2) for i in range(3)}

        return result, confidences

    except Exception as e:
        return f"Prediction error: {str(e)}", {}

# Gradio Interface (updated for newer Gradio)
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload Oral Image"),
    outputs=[
        gr.Label(label="Predicted Condition"),
        gr.JSON(label="Confidence Scores")
    ],
    title="🦷 OralScan AI - Oral Lesion Classifier",
    description="Upload an image to detect oral white lesions using MobileNetV2",
    examples=None,
    flagging_mode="never"   # Updated parameter name
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)