app.py

import streamlit as st
import torch
import torchvision
import torchmetrics
import pytorch_lightning as pl
import numpy as np
import cv2
import time
import pydicom
import nibabel as nib
import io
from torchvision import transforms
from PIL import Image

# Load the trained model
class PneumoniaModel(pl.LightningModule):
    def __init__(self):
        super(PneumoniaModel, self).__init__()
        self.model = torchvision.models.resnet18()
        self.model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        self.model.fc = torch.nn.Linear(in_features=512, out_features=1, bias=True)

        self.loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([5.0]))
        self.val_acc = torchmetrics.Accuracy(task="binary")
        self.val_auc = torchmetrics.AUROC(task="binary")
        self.val_outputs = []

    def forward(self, data):
        return self.model(data)

    def validation_step(self, batch, batch_idx):
        x_ray, label = batch
        label = label.float()
        pred = self(x_ray)[:, 0]
        loss = self.loss_fn(pred, label)
        self.val_outputs.append({"preds": pred, "targets": label})
        return loss

    def on_validation_epoch_end(self):
        all_preds = torch.cat([x["preds"] for x in self.val_outputs]).cpu().numpy()
        all_targets = torch.cat([x["targets"] for x in self.val_outputs]).cpu().numpy()
        self.val_outputs.clear()

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=1e-4)

# Load trained model weights
model = PneumoniaModel()
checkpoint = torch.load("weights_3.ckpt", map_location=torch.device('cpu'), weights_only=False)
state_dict = checkpoint["state_dict"]
model.load_state_dict(state_dict)
model.eval()

# Preprocessing function
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5], std=[0.5])
    ])
    return transform(image).unsqueeze(0)

# Function to load and preprocess different file types
def load_image(file_path, file_type):
    file_type = file_type.lower()
    
    try:
        if file_type in ["png", "jpg", "jpeg"]:
            # For file objects from streamlit
            if hasattr(file_path, 'read'):
                image = Image.open(file_path).convert("L")  # Convert to grayscale
            else:
                image = Image.open(file_path).convert("L")
            image = np.array(image)
            
        elif file_type == "dcm":
            # For file objects from streamlit
            if hasattr(file_path, 'read'):
                # Create a temporary BytesIO object
                temp_file = io.BytesIO(file_path.read())
                file_path.seek(0)  # Reset pointer for future reads
                dicom_data = pydicom.dcmread(temp_file)
            else:
                dicom_data = pydicom.dcmread(file_path)
                
            image = dicom_data.pixel_array
            
        elif file_type in ["nii", "nii.gz"]:
            # For file objects from streamlit
            if hasattr(file_path, 'read'):
                # We need to save temporarily for nibabel
                with open("temp_file." + file_type, "wb") as f:
                    f.write(file_path.read())
                file_path.seek(0)  # Reset pointer for future reads
                nifti_data = nib.load("temp_file." + file_type)
                # Clean up the temp file
                import os
                try:
                    os.remove("temp_file." + file_type)
                except:
                    pass  # Ignore cleanup errors
            else:
                nifti_data = nib.load(file_path)
                
            image = nifti_data.get_fdata()
            image = np.squeeze(image)  # Only one squeeze needed
            
        else:
            return None
        
        # Common processing for all image types
        # Normalize to 0-255 range if needed
        if image.max() > 1.0 and image.max() <= 255:
            # Already in 0-255 range, no need to normalize
            pass
        else:
            # Normalize to 0-255
            image = np.uint8(255 * (image - np.min(image)) / (np.max(image) - np.min(image) + 1e-10))  # Added small value to prevent division by zero
        
        # Resize to model's expected input size
        image = cv2.resize(image, (256, 256))
        
        # Apply the preprocessing and return tensor
        return preprocess_image(image)
        
    except Exception as e:
        import traceback
        st.error(f"Error processing image: {str(e)}")
        st.code(traceback.format_exc())
        return None

# Streamlit Web App
st.set_page_config(
    page_title="PneumoFind",
    page_icon="🫁",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        color: #3498db;
        text-align: center;
        margin-bottom: 1rem;
        font-weight: 700;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #2c3e50;
        text-align: center;
        margin-bottom: 2rem;
    }
    .result-normal {
        padding: 20px;
        border-radius: 10px;
        background-color: #2ecc71;
        color: white;
        text-align: center;
        font-size: 2rem;
        font-weight: bold;
        margin: 20px 0;
    }
    .result-pneumonia {
        padding: 20px;
        border-radius: 10px;
        background-color: #e74c3c;
        color: white;
        text-align: center;
        font-size: 2rem;
        font-weight: bold;
        margin: 20px 0;
    }
    .upload-section {
        background-color: #f8f9fa;
        padding: 30px;
        border-radius: 15px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        margin-bottom: 30px;
    }
    .footer {
        text-align: center;
        color: #7f8c8d;
        font-size: 0.9rem;
        padding: 20px;
        border-top: 1px solid #eee;
        margin-top: 40px;
    }
    .stImage img {
        border-radius: 10px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
</style>
""", unsafe_allow_html=True)

# Header
st.markdown("<h1 class='main-header'>PneumoFind</h1>", unsafe_allow_html=True)
st.markdown("<h2 class='sub-header'>Advanced AI-Powered Pneumonia Detection</h2>", unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.image("steth.png")
    st.markdown("## About PneumoFind")
    st.info(
        "PneumoFind uses deep learning to analyze chest X-rays and "
        "detect signs of pneumonia with high accuracy. Upload your medical "
        "image for instant analysis."
    )
    
    st.markdown("## Supported Formats")
    st.markdown("- X-ray Images (PNG, JPG, JPEG)")
    
    
    st.markdown("## Interpretation")
    st.success("**Normal**: No signs of pneumonia detected")
    st.error("**Pneumonia**: Signs of pneumonia detected")

# Main content
st.markdown("<div class='upload-section'>", unsafe_allow_html=True)
uploaded_file = st.file_uploader("Upload an X-ray image for analysis", type=["png", "jpg", "jpeg"])
st.markdown("</div>", unsafe_allow_html=True)

# Process image if uploaded
if uploaded_file is not None:
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### Uploaded Image")
        st.image(uploaded_file, caption="X-ray Image", use_container_width=True)
    
    with col2:
        st.markdown("### Analysis Results")
        
        # Progress bar for analysis simulation
        with st.spinner("Analyzing image..."):
            # Process the image
            file_extension = uploaded_file.name.split(".")[-1]
            processed_image = load_image(uploaded_file, file_extension)
            
            if processed_image is not None:
                # Process with model
                progress_bar = st.progress(0)
                for i in range(100):
                    time.sleep(0.01)  # Add a small delay for visual effect
                    progress_bar.progress(i + 1)

                # Get prediction
                with torch.no_grad():
                    output = model(processed_image)  # Model outputs raw logits
                    probability = torch.sigmoid(output).item()  # Convert logits to probability
                    prediction = "Pneumonia Detected" if probability > 0.15 else "No Pneumonia Detected"

                # Display results
                if probability > 0.15:
                    st.markdown(f"<div class='result-pneumonia'>{prediction}</div>", unsafe_allow_html=True)
                    st.warning(f"Confidence Score: {probability:.2f}")  # Display correct probability
                    st.markdown("#### Recommendation")
                    st.error("Please consult a healthcare professional for proper diagnosis and treatment.")
                else:
                    st.markdown(f"<div class='result-normal'>{prediction}</div>", unsafe_allow_html=True)
                    st.info(f"Confidence Score: {1 - probability:.2f}")  # Correct confidence display
                    st.markdown("#### Recommendation")
                    st.success("X-ray appears normal. Continue regular health monitoring.")

            else:
                st.error("Error: File format not supported or corrupted image.")
else:
    # Display sample image gallery when no file is uploaded
    st.markdown("### Sample X-rays")
    st.info("Upload an X-ray image to get started. Here are example images for reference.")
    
    sample_col1, sample_col2 = st.columns(2)
    with sample_col1:
        st.image("nopneumoniaxray.png", 
                caption="Example of a normal chest X-ray", width=300)
    with sample_col2:
        st.image("pneumoniaxray.png", 
                caption="Example of a pneumonia chest X-ray", width=300)

# Informational section
st.markdown("## About Pneumonia")
expander = st.expander("Learn more about pneumonia")
with expander:
    st.markdown("""
    Pneumonia is an infection that inflames the air sacs in one or both lungs. The air sacs may fill with fluid or pus, 
    causing cough with phlegm or pus, fever, chills, and difficulty breathing. Various organisms, including bacteria, 
    viruses and fungi, can cause pneumonia.
    
    **Common symptoms include:**
    - Chest pain when breathing or coughing
    - Confusion or changes in mental awareness (in adults age 65 and older)
    - Cough, which may produce phlegm
    - Fatigue
    - Fever, sweating and shaking chills
    - Lower than normal body temperature (in adults older than age 65 and people with weak immune systems)
    - Nausea, vomiting or diarrhea
    - Shortness of breath
    """)

# Footer
st.markdown("<div class='footer'>App Developed by Syed Faizan | © 2025 PneumoFind</div>", unsafe_allow_html=True)