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

MODEL_NAME = "pneumonia_detection_model.keras"

HF_REPO_ID = "saad1BM/pneumonia-detection-system"  
CLASS_NAMES = ['NORMAL', 'PNEUMONIA']
IMAGE_SIZE = (224, 224)

@st.cache_resource
def load_pneumonia_model():
  
    model_path = hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_NAME)
  
    model = load_model(model_path)
    return model


st.set_page_config(page_title="Pneumonia Detection System (AI Powered)", layout="centered")

st.title("🫁 Pneumonia Detection System (AI Powered)")
st.caption("Upload a chest X-ray image to predict Normal or Pneumonia.")


try:
    model = load_pneumonia_model()
except Exception as e:
    st.error(f"Error loading model from Hugging Face: {e}")
    st.info("Please make sure the model file is correctly uploaded to Hugging Face Hub.")
    st.stop()


uploaded_file = st.file_uploader("Choose a Chest X-ray Image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
   
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption='Uploaded X-ray Image', use_column_width=True)
    
  
    if st.button("Detect Pneumonia"):
        st.subheader("📊 Prediction Result")
        
        with st.spinner('Analyzing X-ray image...'):
          
            img = image.resize(IMAGE_SIZE)
            img_array = np.array(img) / 255.0 
            img_array = np.expand_dims(img_array, axis=0) # Batch dimension add karna
            
         
            predictions = model.predict(img_array)
            score = tf.nn.softmax(predictions[0])
            
            
            predicted_class_index = np.argmax(score)
            predicted_class = CLASS_NAMES[predicted_class_index]
            confidence = np.max(score) * 100

            if predicted_class == 'PNEUMONIA':
                st.error(f"### ⚠️ Prediction: {predicted_class}")
                st.markdown(f"**Confidence:** **{confidence:.2f}%**")
                
            else:
                st.success(f"### ✅ Prediction: {predicted_class}")
                st.markdown(f"**Confidence:** **{confidence:.2f}%**")

           
            st.markdown("---")
            st.bar_chart({
                "Normal": score[0].numpy(),
                "Pneumonia": score[1].numpy()
            })