import streamlit as st
import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import load_model

# Define constants
IMG_WIDTH, IMG_HEIGHT = 224, 224  # Adjust if your model requires different dimensions
model_path = 'dog_cat_classifier_model.keras'  # Replace with your model path

# Load the model (outside the main app function for efficiency)
model = load_model(model_path)

def main():
    """
    Streamlit app for image classification with user-friendly interface.
    """

    # Title and description
    st.title("Intriguing Image Classifier")
    st.write("Upload an image and discover its most likely category along with probabilities in a compelling way!")

    # File uploader and sidebar for image selection
    uploaded_file = st.file_uploader("Choose an Image", type=['jpg', 'jpeg', 'png'])
    image_selected = st.sidebar.selectbox("Select Image", (None, "Uploaded Image"))

    if uploaded_file is not None:
        image_display = image.load_img(uploaded_file, target_size=(IMG_WIDTH, IMG_HEIGHT))
        st.image(image_display, caption="Uploaded Image", use_column_width=True)
        image_selected = "Uploaded Image"

    # Preprocess image if one is selected
    if image_selected:
        img_array = image.img_to_array(image_display)
        img_array = np.expand_dims(img_array, axis=0)
        img_array /= 255.0  # Rescale pixel values to [0, 1]

        # Make predictions and get class labels (assuming your model outputs probabilities)
        predictions = model.predict(img_array)
        class_labels = [f"{label}: {prob:.2%}" for label, prob in zip(get_class_labels(model), predictions[0])]

        # Display predictions in an intriguing way (replace with your preferred method)
        st.header("Predictions:")
        progress_bar_width = 800  # Adjust for desired visual style

        for label, prob in zip(class_labels, predictions[0]):
            progress_bar = st.progress(label)
            progress_bar.progress(int(prob * 100))  # Update progress bar based on probability

# Function to retrieve class labels from the model (replace if your model structure is different)
def get_class_labels(model):
    class_names = list(model.class_names)  # Assuming class names are directly accessible
    return class_names

main()

# if __name__ == '__main__':
#     main()