src/streamlit_app.py

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

@keras.saving.register_keras_serializable(package="CustomLayers")
class GrayscaleToRGB(tf.keras.layers.Layer):
    def call(self, inputs):
        return tf.image.grayscale_to_rgb(inputs)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1], input_shape[2], 3)

MODEL_PATH = "src/best_model.keras"
IMG_SIZE = (160, 160)
THRESHOLD = 0.5

@st.cache_resource
def load_model():
    return tf.keras.models.load_model(
        MODEL_PATH,
        compile=False,
        safe_mode=False,
        custom_objects={
            "GrayscaleToRGB": GrayscaleToRGB,
            "CustomLayers>GrayscaleToRGB": GrayscaleToRGB
        }
    )

model = load_model()

def preprocess(pil_img):
    img = pil_img.convert("L").resize(IMG_SIZE)
    arr = np.array(img).astype("float32") / 255.0
    arr = np.expand_dims(arr, axis=-1)
    arr = np.expand_dims(arr, axis=0)
    return arr

st.title("Pneumonia Detection App")
st.write("Upload a chest X-ray image and get predicted class + probability.")
st.markdown(
    '<h5><span style="color:red">If you have DICOM (.dcm) files convert them to PNG or JPG and then upload</span>.</h5>',
    unsafe_allow_html=True
)

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

if uploaded:
    img = Image.open(uploaded)
    st.image(img, caption="Uploaded image", use_container_width=True)

    x = preprocess(img)
    prob = float(model.predict(x, verbose=0).ravel()[0])
    pred = "Pneumonia" if prob >= THRESHOLD else "No Pneumonia"

    st.subheader("Prediction")
    st.write(f"**Class:** {pred}")
    st.write(f"**Probability (Pneumonia):** {prob:.4f}")