import streamlit as st
import os
from PIL import Image
import numpy as np
import cv2
from keras.models import load_model
from util import set_background  


st.set_page_config(page_title="Pneumonia Classifier", layout="centered")
st.title("🩺 Pneumonia Classifier Application")
st.header("Upload a Chest X-ray Image")


@st.cache_resource
def load_bg():
    bg_path = os.path.join(os.path.dirname(__file__), "bgs", "bg5.png")
    bg = cv2.imread(bg_path)
    if bg is None:
        return None
    blurred = cv2.GaussianBlur(bg, (15, 15), 0)
    return blurred

@st.cache_resource
def load_pneumonia_model():
    model_path = os.path.join(
        os.path.dirname(__file__),
        "model",
        "pneumonia_classifier.keras"
    )
    return load_model(model_path)


bg_img = load_bg()
if bg_img is not None:
    tmp_bg = "/tmp/bg_blur.png"
    cv2.imwrite(tmp_bg, bg_img)
    set_background(tmp_bg)


model = load_pneumonia_model()
class_names = ["NORMAL", "PNEUMONIA"]


file = st.file_uploader(
    "Upload a chest X-ray image",
    type=["jpeg", "jpg", "png"]
)


if file is not None:
    image = Image.open(file).convert("L")
    st.image(image, caption="Uploaded X-ray", use_container_width=True)

    img = np.array(image)
    img = cv2.resize(img, (128, 128))
    img = img / 255.0
    img = np.expand_dims(img, axis=(0, -1))

    prediction = model.predict(img)
    class_idx = np.argmax(prediction)
    confidence = np.max(prediction)

    st.success(f"Prediction: {class_names[class_idx]}")
    st.info(f"Confidence: {confidence * 100:.2f}%")