import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image

# Load the model
model = tf.keras.models.load_model("vgg19_binary_nonbinary.h5")

def preprocess_image(image):
    # Convert RGBA to RGB if the image has an alpha channel
    if image.mode == "RGBA":
        image = image.convert("RGB")
    # Resize and normalize the image
    image = image.resize((224, 224))  # Resize to match model input size
    image = np.array(image) / 255.0   # Normalize pixel values
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    return image
    
# Streamlit app
st.title("Binary vs Non-Binary Image Classification")
st.write("Upload an image to classify it as 'binary' or 'non-binary'.")

# File uploader
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.write("Classifying...")

    # Preprocess and predict
    processed_image = preprocess_image(image)
    predictions = model.predict(processed_image)
    class_names = ["binary", "non-binary"]
    confidence = {class_names[i]: float(predictions[0][i]) for i in range(2)}

    # Display the prediction
    st.write("Prediction:")
    st.write(confidence)