import streamlit as st
import numpy as np
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.densenet import preprocess_input
from tensorflow.keras.models import load_model

# Load the pre-trained model
model = load_model("./best_weights.hdf5")  # Replace with the path to your model

class_labels = {0: 'COVID19', 1: 'NORMAL', 2: 'PNEUMONIA', 3: 'TURBERCULOSIS'}

def main():
    st.title("Respiratory Symptom Prediction")
    st.write("Upload an image to predict the respiratory symptom.")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        # Load and preprocess the uploaded image
        image = load_img(uploaded_file, target_size=(224, 224))
        image_array = img_to_array(image)
        image_preprocessed = preprocess_input(image_array)
        image_batch = np.expand_dims(image_preprocessed, axis=0)

        # Make predictions
        predictions = model.predict(image_batch)
        st.write(predictions)
        predicted_class_index = np.argmax(predictions)
        predicted_class_label = class_labels[predicted_class_index]

        # Display the uploaded image and prediction result
        st.image(image, caption="Uploaded Image", use_column_width=True)
        st.write(f"Predicted class label: {predicted_class_label}")

if __name__ == "__main__":
    main()