import numpy as np
import streamlit as st
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.resnet50 import preprocess_input
import matplotlib.pyplot as plt

# Load the trained model
model_path = 'my_cnn.h5'  # or '/content/my_model.keras'
model = load_model(model_path)

# Preprocess the image
def preprocess_image(img):
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    img_array = preprocess_input(img_array)  # Ensure correct preprocessing for ResNet50
    return img_array

# Make predictions and map to class labels
def classify_image(img):
    img_array = preprocess_image(img)
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions, axis=1)  # Get the index of the highest probability

    class_labels = {0: 'Aedes Aegypti', 1: 'Anopheles Stephensi', 2: 'Culex Quinquefasciatus'}
    species = class_labels.get(predicted_class[0], "Unknown")
    
    return species, predictions

# Streamlit application
def main():
    st.title("Mosquito Species Classification")
    st.write("Upload a mosquito image to classify its species.")

    # File uploader for image input
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        # Load the image for display
        img = image.load_img(uploaded_file, target_size=(224, 224))
        st.image(img, caption='Uploaded Image', use_column_width=True)
        
        # Classify the image
        result, probabilities = classify_image(img)
        st.write(f'Predicted mosquito species: **{result}**')
        st.write(f'Prediction probabilities: {probabilities}')

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