import gradio as gr
from tensorflow import keras
import numpy as np
from PIL import Image
import cv2
import tensorflow as tf
import base64
import io

model = keras.models.load_model('my_model (2).h5')

CLASS_NAMES = ['Non-Tumor', 'Non-Viable-Tumor', 'Viable', 'viable: non-viable']

def make_gradcam_heatmap(img_array, model):
    last_conv_layer = model.get_layer('last_conv_layer')
    grad_model = tf.keras.models.Model(
        [model.inputs],
        [last_conv_layer.output, model.output]
    )
    
    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(img_array)
        pred_index = tf.argmax(predictions[0])
        class_channel = predictions[:, pred_index]
    
    grads = tape.gradient(class_channel, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    
    conv_outputs = conv_outputs[0]
    heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)
    heatmap = tf.maximum(heatmap, 0) / (tf.math.reduce_max(heatmap) + 1e-8)
    
    return heatmap.numpy()

def predict_from_base64(base64_string):
    try:
        if ',' in base64_string:
            base64_string = base64_string.split(',')[1]
        
        image_bytes = base64.b64decode(base64_string)
        img = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        
        img = img.resize((224, 224))
        img_array = np.array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        
        predictions = model.predict(img_array)
        pred_class = CLASS_NAMES[np.argmax(predictions[0])]
        confidence = float(np.max(predictions[0])) * 100
        
        result_text = f"Prediction: {pred_class}\nConfidence: {confidence:.2f}%\n\n"
        result_text += "All Probabilities:\n"
        for i, name in enumerate(CLASS_NAMES):
            result_text += f"  {name}: {predictions[0][i]*100:.2f}%\n"
        
        try:
            heatmap = make_gradcam_heatmap(img_array, model)
            heatmap = cv2.resize(heatmap, (224, 224))
            heatmap = np.uint8(255 * heatmap)
            heatmap_colored = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
            original = np.array(img)
            superimposed = cv2.addWeighted(original, 0.6, heatmap_colored, 0.4, 0)
            output_image = superimposed
        except Exception as e:
            output_image = np.array(img)
            result_text += f"\n(Grad-CAM unavailable: {str(e)})"
        
        return output_image, result_text
    
    except Exception as e:
        return None, f"Error: {str(e)}"

def predict_from_image(input_image):
    img = Image.fromarray(input_image).convert('RGB')
    img = img.resize((224, 224))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    
    predictions = model.predict(img_array)
    pred_class = CLASS_NAMES[np.argmax(predictions[0])]
    confidence = float(np.max(predictions[0])) * 100
    
    result_text = f"Prediction: {pred_class}\nConfidence: {confidence:.2f}%\n\n"
    result_text += "All Probabilities:\n"
    for i, name in enumerate(CLASS_NAMES):
        result_text += f"  {name}: {predictions[0][i]*100:.2f}%\n"
    
    try:
        heatmap = make_gradcam_heatmap(img_array, model)
        heatmap = cv2.resize(heatmap, (224, 224))
        heatmap = np.uint8(255 * heatmap)
        heatmap_colored = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        original = np.array(img)
        superimposed = cv2.addWeighted(original, 0.6, heatmap_colored, 0.4, 0)
        output_image = superimposed
    except Exception as e:
        output_image = np.array(img)
        result_text += f"\n(Grad-CAM unavailable: {str(e)})"
    
    return output_image, result_text

image_interface = gr.Interface(
    fn=predict_from_image,
    inputs=gr.Image(label="Upload Histopathology Image"),
    outputs=[
        gr.Image(label="Grad-CAM Visualization"),
        gr.Textbox(label="Classification Result")
    ],
    title="Bone Cancer Detection (Osteosarcoma)",
    description="Upload an H&E stained histopathology image."
)

api_interface = gr.Interface(
    fn=predict_from_base64,
    inputs=gr.Textbox(label="Base64 Image String"),
    outputs=[
        gr.Image(label="Grad-CAM Visualization"),
        gr.Textbox(label="Classification Result")
    ],
    api_name="predict_base64"
)

demo = gr.TabbedInterface(
    [image_interface, api_interface],
    ["Upload Image", "API (Base64)"]
)

demo.launch()