import streamlit as st
import numpy as np
from PIL import Image
from keras.models import load_model

# Load the pre-trained model for banana ripeness detection
banana_model = load_model("trained model/best_model.h5")

# Define class names for the banana disease detection
class_names_disease = {
    0: 'BUNCHY_TOP',
    1: 'CORDANA',
    2: 'PANAMA',
    3: 'SIGATOKA'
}

# Define class names for the banana ripeness detection
class_names_ripeness = ["Banana_G1", "Banana_G2", "Rotten"]
model = load_model("trained model/best_model.h5")

def preprocess_image(image):
    img = Image.open(image)
    img = img.resize((256, 256))  # Resize the image to the input size of the model
    img_array = np.array(img)
    img_array = img_array / 255.0  # Normalize the pixel values
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    return img_array


def predict(image):
    img_array = preprocess_image(image)
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions)
    predicted_label = class_names_disease[predicted_class]
    return predicted_label

def predict_disease(uploaded_file):
    if uploaded_file is not None:
        predicted_label = predict(uploaded_file)
        return predicted_label


def predict_ripeness(image):
    img_array = preprocess_image(image)
    predictions = banana_model.predict(img_array)
    predicted_class = np.argmax(predictions)
    predicted_label = class_names_ripeness[predicted_class]
    return predicted_label

def main():
    st.title("Banana Analysis App")
    st.write("Choose an option to analyze bananas")

    # Options for banana analysis
    analysis_option = st.radio("Choose an option", ["Banana Disease Detection", "Banana Ripeness Detection"])

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

    if uploaded_file is not None:
        # Display the uploaded image
        st.image(uploaded_file, caption='Uploaded Image', use_column_width=True)
        if st.button("Analyze"):
            if analysis_option == "Banana Disease Detection":
                predicted_label = predict_disease(uploaded_file)
                st.success(f"Predicted disease: {predicted_label}")
            elif analysis_option == "Banana Ripeness Detection":
                predicted_label = predict_ripeness(uploaded_file)
                st.success(f"Predicted ripeness: {predicted_label}")

if __name__ == '__main__':
    main()