import cv2
import os
import PIL.Image
import numpy as np
import streamlit as st
import tensorflow as tf
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2,preprocess_input as mobilenet_v2_preprocess_input

model = tf.keras.models.load_model("model_cassava.hdf5")
### load file
# uploaded_file = st.file_uploader("Choose a image file", type="jpg")

map_dict = {0: 'Cassava Bacterial Blight (CBB)',
            1: 'Cassava Brown Streak Disease (CBSD)',
            2: 'Cassava Green Mottle (CGM)',
            3: 'Cassava Mosaic Disease (CMD)',
            4: 'Healthy'}

option = st.radio('', ['Choose a test image', 'Choose your own image'])
if option == 'Choose your own image':
    uploaded_file = st.file_uploader("Choose a image file", type="jpg")

    if uploaded_file is not None:
        # Convert the file to an opencv image.
        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
        opencv_image = cv2.imdecode(file_bytes, 1)
        opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(opencv_image,(224,224))
        # Now do something with the image! For example, let's display it:
        st.image(opencv_image, channels="RGB")

        # resized = mobilenet_v2_preprocess_input(resized)
        img_reshape = resized[np.newaxis,...]

        Genrate_pred = st.button("Generate Prediction")
        if Genrate_pred:
            predictions_proba = model.predict(img_reshape).max()*100
            prediction = model.predict(img_reshape).argmax()
            st.title("Predicted Label for the image is {}".format(map_dict[prediction]))
            st.title("with a probability of" f" { '{:.2f}'.format(predictions_proba)}")

else:
    test_images = os.listdir('train_images')
    test_image = st.selectbox('Please select a test image:', test_images)
    file_path = 'train_images/' + test_image
    with open(file_path, 'rb') as img_stream:
        file_bytes = np.asarray(bytearray(img_stream.read()), dtype=np.uint8)
        opencv_image = cv2.imdecode(file_bytes, 1)
        opencv_image = cv2.cvtColor(opencv_image, cv2.COLOR_BGR2RGB)
        resized = cv2.resize(opencv_image,(224,224))
        # Now do something with the image! For example, let's display it:
        st.image(opencv_image, channels="RGB")

        # resized = mobilenet_v2_preprocess_input(resized)
        img_reshape = resized[np.newaxis,...]

        Genrate_pred = st.button("Generate Prediction")
        if Genrate_pred:
            predictions_proba = model.predict(img_reshape).max()*100
            prediction = model.predict(img_reshape).argmax()
            st.title("Predicted Label for the image is {}".format(map_dict[prediction]))
            st.title("with a probability of" f" { '{:.2f}'.format(predictions_proba)}")