import os
os.environ["TF_USE_LEGACY_KERAS"] = "1"

from keras.models import load_model
import streamlit as st
import numpy as np
from io import BytesIO
from PIL import Image
import tensorflow as tf

st.markdown(
    """
    <style>
    .reportview-container {
        background: url('./bg.jpg');
        background-size: cover;
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.markdown("# Bananas Maturity Classification ")
st.sidebar.markdown("# Main Page")

MODEL = load_model("./1")

CLASS_NAMES = ["Banana_G1", "Banana_G2", "Rotten"]


def read_file_as_image(data) -> np.ndarray:
    image = np.array(Image.open(BytesIO(data)))
    return image


def predict(
    file,
):
    image = read_file_as_image(file.read())
    shape = image.shape
    img_batch = np.expand_dims(image, 0)
    # resize image to (256,256,3)
    img_batch = tf.image.resize(img_batch, (256, 256))
    prediction = MODEL.predict(img_batch)
    predicted_class = CLASS_NAMES[np.argmax(prediction[0])]
    confidence = np.max(prediction[0])
    if predicted_class == "Banana_G2":
        predicted_class = "Green Banana- not ripen"
    elif predicted_class == "Banana_G1":
        predicted_class = "Mature Banana -ripen"
    else:
        predicted_class = "Rotten Banana"
    return {
        'class': predicted_class,
        'confidence': float(confidence)
    }


st.write("Upload an image or capture one with your camera")

option = st.selectbox("Choose an option", ["Upload Image", "Capture Image"])

if option == "Upload Image":
    uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
    if uploaded_file is not None:
        result = predict(uploaded_file)
        predicted_class = result['class']
        confidence = result['confidence']
        if predicted_class == "Green Banana- not ripen":
            color = 'green'
        elif predicted_class == "Mature Banana -ripen":
            color = 'yellow'
        else:
            color = 'red'
        st.markdown(
            f'<p style="color:{color}; font-size:24px;">Predicted class: {predicted_class}, Confidence: {confidence:.2f}</p>',
            unsafe_allow_html=True)