import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
import matplotlib.pyplot as plt
from tensorflow.keras.models import Model

# Load trained model
model = tf.keras.models.load_model("densenet_model.h5")
resize_dim = (256, 256)
classes = ["OK", "Defective"]

# Preprocessing function
def preprocess_image(img):
    if img.ndim == 3 and img.shape[2] == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    elif img.ndim == 2:
        gray = img
    else:
        raise ValueError("Unsupported image shape")
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    cl_img = clahe.apply(gray)
    denoised = cv2.fastNlMeansDenoising(cl_img, None, h=10, templateWindowSize=7, searchWindowSize=21)
    resized = cv2.resize(denoised, resize_dim)
    sobel_x = cv2.Sobel(resized, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(resized, cv2.CV_64F, 0, 1, ksize=3)
    sobel = np.sqrt(sobel_x**2 + sobel_y**2)
    sobel = np.uint8(np.clip(sobel, 0, 255))
    normalized = sobel / 255.0
    normalized = normalized[..., np.newaxis]
    rgb_img = np.repeat(normalized, 3, axis=-1)
    return np.expand_dims(rgb_img, axis=0).astype(np.float32), rgb_img  # Return both for Grad-CAM

# Grad-CAM heatmap generation
def get_gradcam_heatmap(model, img_array, layer_name='conv5_block32_concat', class_index=None):
    grad_model = Model(inputs=model.input,
                       outputs=[model.get_layer(layer_name).output, model.output])
    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(img_array)
        if class_index is None:
            class_index = tf.argmax(predictions[0])
        loss = predictions[:, class_index]
    grads = tape.gradient(loss, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    conv_outputs = conv_outputs[0]
    heatmap = tf.reduce_sum(tf.multiply(pooled_grads, conv_outputs), axis=-1)
    heatmap = np.maximum(heatmap, 0)
    heatmap /= tf.reduce_max(heatmap)
    return heatmap.numpy()

# Overlay heatmap on image
def overlay_heatmap(img, heatmap, alpha=0.4):
    heatmap_resized = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
    heatmap_colored = cv2.applyColorMap(np.uint8(255 * heatmap_resized), cv2.COLORMAP_JET)
    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
    overlay = np.uint8(heatmap_colored * alpha + img * 255)
    overlay = np.clip(overlay, 0, 255)
    return overlay.astype(np.uint8)

# Combined prediction + Grad-CAM function
def predict_with_gradcam(img):
    try:
        processed_img, img_rgb = preprocess_image(img)
        preds = model.predict(processed_img)[0]
        class_idx = np.argmax(preds)
        confidence = preds[class_idx]

        heatmap = get_gradcam_heatmap(model, processed_img, class_index=class_idx)
        overlay_img = overlay_heatmap(img_rgb, heatmap)

        label = f"Prediction: {classes[class_idx]} (Confidence: {confidence:.2f})"
        return overlay_img, label
    except Exception as e:
        return np.zeros((256, 256, 3), dtype=np.uint8), f"Error: {str(e)}"

# Gradio interface
iface = gr.Interface(
    fn=predict_with_gradcam,
    inputs=gr.Image(type="numpy", label="Upload Casting Image (Grayscale or RGB)"),
    outputs=[
        gr.Image(type="numpy", label="Grad-CAM Heatmap Overlay"),
        gr.Textbox(label="Prediction")
    ],
    title="Casting Defect Detection with Grad-CAM",
    description="Upload a casting product image. The model classifies it and highlights image regions influencing its decision."
)

if __name__ == "__main__":
    iface.launch()