import streamlit as st
import numpy as np
from PIL import Image
import tensorflow as tf
from huggingface_hub import hf_hub_download

st.set_page_config(page_title="Brain Tumor AI", page_icon="🧠", layout="wide")
st.title("🧠 Brain Tumor AI ")

# --- Load model from Hugging Face Hub ---
REPO_ID = "hellosara/brain-tumor-mri-model"  # your model repo
FILE_NAME = "1.keras"
CLASS_NAMES = ["Brain Tumor", "Healthy"]

@st.cache_resource(show_spinner=True)
def load_model():
    # Download the model from HF Hub
    model_path = hf_hub_download(repo_id=REPO_ID, filename=FILE_NAME)
    return tf.keras.models.load_model(model_path, compile=False)

MODEL = load_model()

# --- File uploader ---
uploaded_file = st.file_uploader("Upload MRI", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    st.image(uploaded_file, caption="Selected Image", width=300)
    
    if st.button("Predict"):
        try:
            # Keep logic simple: just convert to RGB and numpy array
            image = Image.open(uploaded_file).convert("RGB")
            img_array = np.array(image)  # no resizing, no normalization
            img_batch = np.expand_dims(img_array, axis=0)
            
            # Predict
            predictions = MODEL.predict(img_batch)
            predicted_class = CLASS_NAMES[np.argmax(predictions[0])]
            confidence = float(np.max(predictions[0]))
            
            st.success(f"Result: {predicted_class}")
            st.write(f"Confidence: {confidence:.4f}")
        except Exception as e:
            st.error(f"Error: {e}")