import streamlit as st
import numpy as np
import tensorflow as tf
import pydicom
from PIL import Image
import io

@st.cache_resource
def load_model():
    return tf.keras.models.load_model("best_cnn_model.h5")

model = load_model()
IMG_SIZE = 128

st.title("Pneumonia Detection from Chest X-ray")
st.write("Upload a chest X-ray image (DICOM format)")

uploaded_file = st.file_uploader("Upload DICOM Image", type=None)

def preprocess_dicom_bytes(dicom_bytes):
    dicom = pydicom.dcmread(io.BytesIO(dicom_bytes))
    img = dicom.pixel_array.astype("float32")

    img = img / np.max(img)

    img = Image.fromarray((img * 255).astype(np.uint8))
    img = img.resize((IMG_SIZE, IMG_SIZE))

    img = np.array(img) / 255.0
    img = img.reshape(1, IMG_SIZE, IMG_SIZE, 1)
    return img

if uploaded_file:
    try:
        image = preprocess_dicom_bytes(uploaded_file.read())
        prob = model.predict(image)[0][0]

        label = "Pneumonia" if prob >= 0.45 else "No Pneumonia"

        st.success("Prediction complete")
        st.write(f"**Predicted Class:** {label}")
        st.write(f"**Probability:** {prob:.2f}")

    except Exception as e:
        st.error(f"Error processing file: {e}")