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

# Load the trained model
model = tf.keras.models.load_model("vgg19_bone_fracture_model.h5")

# Preprocess uploaded image (matching your training pipeline)
def preprocess_image(image, target_size=(128, 128)):
    # Convert image to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    
    # Resize image to target size
    resized = cv2.resize(gray, target_size)
    
    # Normalize image
    normalized = resized / 255.0

    # Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    img_clahe = clahe.apply((normalized * 255).astype(np.uint8))

    # Apply Otsu thresholding
    _, img_otsu = cv2.threshold(img_clahe, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # Apply Canny edge detection
    edges = cv2.Canny(img_otsu, 50, 150)

    # Return the final image with 3 channels (for CNN input)
    stacked = np.stack([normalized] * 3, axis=-1)
    return np.expand_dims(stacked, axis=0)

# Prediction function
def predict(image):
    input_data = preprocess_image(image)
    prediction = model.predict(input_data)[0]
    class_names = ['Not Fractured', 'Fractured']
    # Return the prediction probabilities for each class
    return {class_names[i]: float(prediction[i]) for i in range(2)}

# Gradio app with submit button and colorful interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Bone X-ray"),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="Bone Fracture Detection (CNN Model)",
    description=(
        "Upload a bone X-ray image to predict whether it shows a fracture or not. "
        "The model will provide the likelihood for each class: 'Fractured' or 'Not Fractured'."
    ),
    theme="compact",  # Compact layout for a cleaner design
    live=False,  # Disable live prediction, only after submit
    allow_flagging="never",  # Disable flagging
    css="""
        .gradio-container {
            background-color: #f0f8ff;  /* Light blue background */
            border-radius: 15px;
            box-shadow: 0 8px 16px rgba(0, 0, 0, 0.2);
        }
        .gradio-title {
            color: #004b8d;  /* Dark blue for title */
            font-family: 'Arial', sans-serif;
        }
        .gradio-description {
            color: #555555;  /* Dark grey for description */
            font-size: 16px;
        }
        .gradio-button {
            background-color: #008cba;  /* Blue button */
            color: white;
            font-weight: bold;
            border-radius: 12px;
        }
        .gradio-button:hover {
            background-color: #006f89;  /* Darker blue on hover */
        }
        .gradio-output {
            font-size: 18px;
            font-weight: bold;
            color: #008cba;  /* Blue color for output */
        }
    """,
)

# Run the app (only for local testing)
if __name__ == "__main__":
    iface.launch()