import gradio as gr
import tensorflow as tf
from tensorflow.keras.applications.inception_resnet_v2 import preprocess_input
from tensorflow.keras.preprocessing import image
import numpy as np
from huggingface_hub import hf_hub_download

# Download model from Hugging Face Space
model_path = hf_hub_download(repo_id="Suphawan/Melanoma", filename="best_model_InceptionV2.keras")

# Load the trained InceptionV2 model
try:
    model = tf.keras.models.load_model(model_path)
    print("Model loaded successfully.")
except OSError as e:
    print(f"Error loading model: {e}")
    model = None  # Handle this case appropriately in your code

# Function for prediction
def predict(img):
    if model is None:
        return "Model could not be loaded."
    
    try:
        img_resized = img.resize((224, 224))  # Resize image to the target size
        img_array = image.img_to_array(img_resized)  # Convert image to array
        img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
        img_array = preprocess_input(img_array)  # Preprocess image according to model requirements

        predictions = model.predict(img_array)
        class_idx = np.argmax(predictions, axis=1)[0]
        class_labels = ['Benign', 'Malignant']  # Update according to your class labels
        class_label = class_labels[class_idx]
        confidence = float(predictions[0][class_idx])

        return f"Class: {class_label}, Confidence: {confidence:.2f}"
    except Exception as e:
        return f"Error: {str(e)}"

# Define the Gradio app
with gr.Blocks() as demo:
    gr.Markdown("Image Classification with InceptionV2")
    
    with gr.Row():
        with gr.Column():
            classify_input = gr.Image(type="pil", label="Upload an Image")
            classify_button = gr.Button("Classify!")
        with gr.Column():
            classify_output = gr.Textbox(label="Classification Result")

    classify_button.click(
        predict,
        inputs=[classify_input],
        outputs=[classify_output]
    )

demo.launch(debug=True)