import streamlit as st
from PIL import Image
import torch
from torchvision import models, transforms

# Load the pre-trained model
model = models.densenet121(pretrained=True)
model.eval()

# Define the image transformations
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])

# Define the class labels
class_labels = ['Normal', 'Pneumonia']

# Create a function to preprocess the image
def preprocess_image(image):
    # Convert the image to RGB
    image = image.convert('RGB')

    # Resize the image to match the model's input shape
    image = image.resize((224, 224))

    # Convert the image to a tensor
    image_tensor = transform(image)

    # Add a batch dimension
    image_tensor = image_tensor.unsqueeze(0)

    return image_tensor

# Create a function to make predictions
def predict(image):
    # Preprocess the image
    preprocessed_image = preprocess_image(image)

    # Make the prediction
    with torch.no_grad():
        output = model(preprocessed_image)
        _, predicted_idx = torch.max(output, 1)
        predicted_label = class_labels[predicted_idx.item()]
    
    return predicted_label

# Create the Streamlit app
def main():
    st.title("Pneumonia Detection")
    st.write("Upload an image and the app will predict if it has pneumonia or not.")
    
    # Upload and display the image
    uploaded_image = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
    
    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        st.image(image, caption="Uploaded Image", use_column_width=True)
        
        # Make a prediction
        predicted_label = predict(image)
        st.write("Prediction:", predicted_label)

# Run the app
if __name__ == '__main__':
    main()