import streamlit as st
import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np

# Load model safely (TF 2.15+)
@st.cache_resource
def load_my_model():
    return tf.keras.models.load_model(
        "src/cactus1.keras",
        compile=False,
        safe_mode=False
    )

model = load_my_model()

def process_image(img):
    img = img.resize((64, 64))
    img = np.array(img)

    # Ensure 3 channels (RGB)
    if img.ndim == 2:  
        img = np.stack([img] * 3, axis=-1)
    elif img.shape[-1] == 4:  # RGBA → RGB
        img = img[:, :, :3]

    img = img / 255.0
    img = np.expand_dims(img, axis=0)  # (1, 64, 64, 3)
    return img

st.title("🌵 Aerial Cactus Identifier")
st.write("This application detects whether an aerial image contains cactus using a Fine-Tuned ResNet50 model.")

file = st.file_uploader("Upload an aerial image", type=["jpg", "jpeg", "png"])

if file is not None:
    img = Image.open(file).convert("RGB")  # Always RGB
    st.image(img, caption="Uploaded Image", use_container_width=True)

    image_tensor = process_image(img)
    prediction = model.predict(image_tensor)[0][0]

    if prediction > 0.5:
        st.success(f"Has Cactus 🌵 (Confidence: {prediction:.2%})")
    else:
        st.error(f"No Cactus 🏜️ (Confidence: {(1 - prediction):.2%})")