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

IMAGE_SIZE = 128
NUM_CLASSES = 3

def preprocess_image(image, target_size, add_clahe=True, clip_limit=4, tile_grid_size=(40,40), all_clahe=True):
    """
    Preprocess the images to remove black borders
    and improve contrast using Y channel and CLAHE
    """
    try:
        # Crop the image to remove black borders
        cropped_rgb_image = remove_black_borders(image)

        if add_clahe:
            equalised_image = apply_clahe(
                cropped_rgb_image, clip_limit, tile_grid_size, all_clahe)

            # Resize the image to target size
            resized_image = cv2.resize(equalised_image, (target_size, target_size))
        else:
            # Resize the image to target size
            resized_image = cv2.resize(cropped_rgb_image, (target_size, target_size))

        return resized_image

    except Exception as e:
        print("Error processing image: ")
        print(e)


def remove_black_borders(image):
    """
    Crop the image to remove black borders
    """
    green_channel_image =  image[:, :, 1]

    # Find the contours in the green channel
    contours, _ = cv2.findContours(
        green_channel_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    # Find the largest contour
    largest_contour = max(contours, key=cv2.contourArea)

    # Get the bounding rectangle of the largest contour
    x, y, w, h = cv2.boundingRect(largest_contour)

    # Create a mask with the same size as the bounding rectangle
    mask = np.zeros((h, w), np.uint8)

    # Draw the largest contour on the mask
    cv2.drawContours(mask, [largest_contour - [x, y]], 0, 255, -1)

    # Convert the mask to a 3 channel image
    mask_3_channel = cv2.merge([mask, mask, mask])

    # Crop the image using the mask
    cropped_image = cv2.bitwise_and(image[y:y+h, x:x+w], mask_3_channel)

    return cropped_image


def apply_clahe(image, clip_limit=4, tile_grid_size=(40,40), all_clahe=False):
    """
    Preprocess the image
    """
    try:
        # Extract the y channel from the cropped rgb image
        yuv_image = cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
        cropped_y_channel_image = yuv_image[:, :, 0]

        # Apply CLAHE to the y channel of the cropped image
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        enhanced_y_channel = clahe.apply(cropped_y_channel_image)

        if all_clahe:
            cropped_u_channel = yuv_image[:, :, 1]
            enhanced_u_channel = clahe.apply(cropped_u_channel)
            cropped_v_channel = yuv_image[:, :, 2]
            enhanced_v_channel = clahe.apply(cropped_v_channel)
            
            enhanced_yuv_image = np.stack([enhanced_y_channel, enhanced_u_channel, enhanced_v_channel], axis=-1)

            # Convert YUV to RGB
            enhanced_rgb_image = cv2.cvtColor(enhanced_yuv_image, cv2.COLOR_YUV2RGB)

            return enhanced_rgb_image
        else:
            # Convert the y channel image to grayscale
            enhanced_grayscale_image = cv2.convertScaleAbs(enhanced_y_channel, alpha=(255/219))
        
            # Repeat the equalised grayscale for all 3-channels
            enhanced_grayscale_3_channels = np.stack((enhanced_grayscale_image,) * 3, axis=-1)

            return enhanced_grayscale_3_channels

    except Exception as e:
        print("Error processing image: ", e)

        
def load_vgg16_model():
    vgg16_bs32_img128_CLAHE4 = tf.keras.models.load_model('./VGG16_32_128_CLAHE_4_40_trainableFalse.h5')
    return vgg16_bs32_img128_CLAHE4

    
def predict_input_image(img):
    # img = preprocess_image(img, IMAGE_SIZE, add_clahe=True, clip_limit=4, tile_grid_size=(40,40), all_clahe=True)
    # img_4d=img.reshape(-1, IMAGE_SIZE, IMAGE_SIZE, 3)
    # img_4d_ensemble = [img_4d for _ in range(len(ensemble_model.input))]
    # prediction=ensemble_model.predict([img_4d_ensemble])[0]

    img = preprocess_image(img, IMAGE_SIZE, add_clahe=True, clip_limit=4, tile_grid_size=(40,40), all_clahe=False)
    img_4d=img.reshape(-1, IMAGE_SIZE, IMAGE_SIZE, 3)
    model = load_vgg16_model()
    prediction=model.predict([img_4d])[0]

    dr_classes = ['No DR', 'Mild DR', 'Referable DR']
    return {dr_classes[i]: float(prediction[i]) for i in range(3)}

image = gr.inputs.Image(shape=(IMAGE_SIZE, IMAGE_SIZE))
label = gr.outputs.Label(num_top_classes=NUM_CLASSES)
iface = gr.Interface(fn=predict_input_image, inputs=image, outputs=label).launch(share=True)
iface.launch()