| 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() | |