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

# Define custom metrics
smooth = 1e-15

def dice_coef(y_true, y_pred):
    y_true = tf.keras.layers.Flatten()(y_true)
    y_pred = tf.keras.layers.Flatten()(y_pred)
    intersection = tf.reduce_sum(y_true * y_pred)
    return (2. * intersection + smooth) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + smooth)

def dice_loss(y_true, y_pred):
    return 1.0 - dice_coef(y_true, y_pred)

# Load the model from the same directory as the script
model_filename = "model.h5"  # Replace with your actual model filename
model_path = os.path.join(os.path.dirname(__file__), model_filename)

def load_model(model_path):
    try:
        model = tf.keras.models.load_model(model_path, custom_objects={'dice_loss': dice_loss, 'dice_coef': dice_coef}, compile=False)
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

# Perform inference
model = load_model(model_path)

def perform_inference(image):
    if model is None:
        print("Model not loaded properly.")
        return None, None, None

    # Preprocess the image
    original_shape = image.shape[:2]    
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_resized = cv2.resize(image, (256, 256))
    image_normalized = image_resized / 255.0
    image_expanded = np.expand_dims(image_normalized, axis=0)

    # Get the mask from the model prediction
    mask = model.predict(image_expanded)[0]
    mask_resized = cv2.resize(mask, (original_shape[1], original_shape[0]))
    mask_binary = (mask_resized > 0.5).astype(np.uint8) * 255

    # Find contours in the binary mask
    contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Apply the mask to the original image (heatmap for visualization)
    heatmap_img = cv2.applyColorMap(mask_binary, cv2.COLORMAP_JET)
    segmented_image = cv2.addWeighted(image, 0.7, heatmap_img, 0.3, 0)
    segmented_image_with_box = segmented_image.copy()

    # Get bounding boxes for all contours and annotate
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        cv2.rectangle(segmented_image_with_box, (x, y), (x + w, y + h), (0, 0, 255), 2)
        cv2.putText(segmented_image_with_box, "Tumour Detected", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)


    #Convert back to BGR
    segmented_image_with_box = cv2.cvtColor(segmented_image_with_box, cv2.COLOR_RGB2BGR)
    return (Image.fromarray(image),
            Image.fromarray(mask_binary.astype(np.uint8)),
            Image.fromarray(segmented_image_with_box))



# Gradio app
def gradio_app():
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="numpy")
                submit_btn = gr.Button("Submit")

            with gr.Column():
                 original_image_output = gr.Image(label="Original Image")
                 mask_output = gr.Image(label="Predicted Mask")
                 segmented_image_output= gr.Image(label="Segmented Image")


        submit_btn.click(perform_inference, inputs=input_image, outputs=[original_image_output, mask_output, segmented_image_output])
    demo.launch()


if __name__ == "__main__":
    gradio_app()