import streamlit as st
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image
import cv2
from PIL import Image
import os

# Load pre-trained model
model = ResNet50(weights='imagenet')
last_conv_layer_name = 'conv5_block3_out'

# Grad-CAM function
def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    grad_model = tf.keras.models.Model(
        [model.inputs],
        [model.get_layer(last_conv_layer_name).output, model.output]
    )
    with tf.GradientTape() as tape:
        conv_outputs, predictions = grad_model(img_array)
        if pred_index is None:
            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)
    return heatmap.numpy()

# Streamlit UI
st.title("🧠 Image Classifier + Grad-CAM Visualizer")

choice = st.selectbox(
    "Select Option:",
    ("Upload Image", "Use Sample Image")
)

if choice == "Upload Image":
    image_= st.file_uploader("📤 Upload an image (JPG, JPEG, PNG)", type=["jpg", "jpeg", "png"])
else:
    image_= Image.open(r"dog.jpg").convert('RGB')

if image_ == None:
    st.write("Upload Image")
else:
# Preprocess image
    image_resized = image_.resize((224, 224))
    x = image.img_to_array(image_resized)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)

# Prediction
    preds = model.predict(x)
    pred_class = decode_predictions(preds, top=1)[0][0]
    st.markdown(f"**Predicted:** `{pred_class[1]}` with `{pred_class[2]*100:.2f}%` confidence")

# Grad-CAM
    heatmap = make_gradcam_heatmap(x, model, last_conv_layer_name)
    img_cv = np.array(image_resized)
    heatmap = cv2.resize(heatmap, (img_cv.shape[1], img_cv.shape[0]))
    heatmap = np.uint8(255 * heatmap)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    superimposed_img = cv2.addWeighted(img_cv, 0.6, heatmap, 0.4, 0)

# Display images
    col1, col2 = st.columns(2)
    with col1:
        st.image(image_resized, caption="Original Image", use_container_width=True)
    with col2:
        st.image(superimposed_img, caption="Grad-CAM", use_container_width=True)

# except Exception as e:
#     st.write("Upload image")