import streamlit as st
from PIL import Image
import matplotlib.pyplot as plt
import tensorflow_hub as hub
import tensorflow as tf
import numpy as np
from tensorflow import keras
from tensorflow.keras.models import *
from tensorflow.keras import preprocessing
import time


st.title('AR-CAD CAD element Classifier')
st.caption('From this demo, your can import any image with CAD graphs, then our model will segmented the detected CAD elements in this window. Please be aware that NOT to send us image file more than 200 MB. For any technique problems, please content ar.cad.connect@gmail.com.')
st.markdown("Our Prediction :")

def main():
    file_uploaded = st.file_uploader("Choose File", type=["png","jpg","jpeg"])
    class_btn = st.button("Classify")
    if file_uploaded is not None:    
        image = Image.open(file_uploaded)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        
    if class_btn:
        if file_uploaded is None:
            st.write("Invalid command, please upload an image")
        else:
            with st.spinner('Model working....'):
 
                predictions = predict(image)

                time.sleep(1)
                st.success('Classified')
                st.write(predictions)


def predict(image):
    model = "wall.tflite"
    interpreter = tf.lite.Interpreter(model_path = model)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    
    input_shape = input_details[0]['shape']
    image = np.array(image.resize((200,200)), dtype=np.float32) 
    
    image = image / 255.0
    image = np.expand_dims(image, axis=0)
    interpreter.set_tensor(input_details[0]['index'], image)
    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])
    probabilities = np.array(output_data[0])
    labels = {1 : "wall"}
    label_to_probabilities = []
    for i, probability in enumerate(probabilities):
        label_to_probabilities.append([labels[i],
float(probability)])
    
    sorted(label_to_probabilities, key=lambda element: element[1])
    result = { 'wall' : 0 }
    
    result = f"{label_to_probabilities[np.argmax(probability)][0]} with a { (100 * np.max(probabilities)).round(2)} % confidence." 
    
    return result

if __name__ == "__main__":
    main()