import streamlit as st
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
import joblib
import gdown

# Google Drive model URLs
KNN_MODEL_URL = 'https://drive.google.com/uc?id=1TJ0KbzFw-2NfuJf67xvp-32uaYLIqpj3'
EXTRACTOR_URL = 'https://drive.google.com/uc?id=1HR2Qc8Fji6RzbtG_K_sqSoiG0AQnvyZa'

# Download the model files
st.write("Downloading models...")
gdown.download(KNN_MODEL_URL, 'knn_pharyngitis_model.pkl', quiet=False)
gdown.download(EXTRACTOR_URL, 'mobilenetv2_feature_extractor.h5', quiet=False)
st.write("Models downloaded successfully!")

# Load the saved models
knn = joblib.load('knn_pharyngitis_model.pkl')
feature_extractor = load_model('mobilenetv2_feature_extractor.h5')

# Function to preprocess the uploaded image
def preprocess_image(image):
    img = image.resize((224, 224))  # Resize to match MobileNetV2 input size
    img_array = np.array(img)
    img_array = preprocess_input(img_array)  # Apply MobileNetV2 preprocessing
    return np.expand_dims(img_array, axis=0)

# Function to classify the image
def classify_image(image):
    processed_image = preprocess_image(image)
    features = feature_extractor.predict(processed_image)
    prediction = knn.predict(features)
    return "Pharyngitis" if prediction[0] == 1 else "No Pharyngitis"

# Streamlit app UI
st.title("Pharyngitis Classification App")
st.write("Upload an image to classify it as 'Pharyngitis' or 'No Pharyngitis'.")

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

if uploaded_file is not None:
    # Load the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Classify the image
    st.write("Classifying...")
    prediction = classify_image(image)
    st.write(f"Prediction: **{prediction}**")