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

# Google Drive file ID
MODEL_URL = "https://drive.google.com/uc?id=1MUaUD7o7euJGWFHDbFo50LZSvP1wsk-y"

# Download model from Google Drive
@st.cache_resource
def download_model():
    output = 'pharyngitis_model.h5'
    gdown.download(MODEL_URL, output, quiet=False)
    model = tf.keras.models.load_model(output)
    return model

# Load model
model = download_model()

# Preprocessing function for the input image
def preprocess_image(image):
    image = Image.open(image).convert('RGB')
    image = image.resize((224, 224))
    image_array = np.array(image) / 255.0  # Normalize the image
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    return image_array

# Prediction function
def predict(image):
    image_array = preprocess_image(image)
    prediction = model.predict(image_array)
    return prediction

# Streamlit UI setup
st.title("Pharyngitis Classifier")
st.write("Upload an image of the throat to check if pharyngitis is detected. Please note that this is not a substitute for professional medical advice. If you feel unwell, consult a doctor.")

# File uploader widget
uploaded_image = st.file_uploader("Upload Image", type=["jpg", "png", "jpeg"])

if uploaded_image is not None:
    # Display a loading animation
    with st.spinner("Please wait... Classifying the image"):
        prediction = predict(uploaded_image)

    # Display result
    if prediction[0] < 0.5:
        st.success("No Pharyngitis Detected")
    else:
        st.error("Pharyngitis Detected")

    # Show uploaded image
    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)