Spaces:
Build error
Build error
| # Import libraries | |
| import gradio as gr | |
| import tensorflow as tf | |
| import numpy as np | |
| # Initialize the number of classes, also the image's height and width | |
| num_classes = 200 | |
| IMG_HEIGHT = 300 | |
| IMG_WIDTH = 300 | |
| # Open the classlabel.txt to read the class labels | |
| with open("classlabel.txt", 'r') as file: | |
| CLASS_LABEL = [x.strip() for x in file.readlines()] | |
| # Function to normalize the image | |
| def normalize_image(img): | |
| img = tf.cast(img, tf.float32)/255. | |
| img = tf.image.resize(img, (IMG_HEIGHT, IMG_WIDTH), method='bilinear') | |
| return img | |
| # Function to select and load the model | |
| def load_model(model_name): | |
| # Load the model based on the model_name input | |
| if model_name == "Xception": | |
| return tf.keras.models.load_model("model/Xception.h5") | |
| elif model_name == "InceptionV3": | |
| return tf.keras.models.load_model("model/InceptionV3.h5") | |
| elif model_name == "InceptionResNetV2": | |
| return tf.keras.models.load_model("model/InceptionResNetV2.h5") | |
| elif model_name == "DenseNet201": | |
| return tf.keras.models.load_model("model/DenseNet201.h5") | |
| else: | |
| raise ValueError("Invalid model_name") | |
| # Main function, let the model make the prediction on the image uploaded | |
| def predict_top_classes(img, model_name): | |
| img = img.convert('RGB') | |
| img_data = normalize_image(img) | |
| x = np.array(img_data) | |
| x = np.expand_dims(x, axis=0) | |
| model = load_model(model_name) | |
| temp = model.predict(x) | |
| idx = np.argsort(np.squeeze(temp))[::-1] | |
| top5_value = np.asarray([temp[0][i] for i in idx[0:5]]) | |
| top5_idx = idx[0:5] | |
| # Return the top 5 highest probability class labels | |
| return {CLASS_LABEL[i]: str(v) for i, v in zip(top5_idx, top5_value)} | |
| # Define the interface | |
| interface = gr.Interface( | |
| predict_top_classes, | |
| [ | |
| gr.Image(type='pil'), | |
| gr.Dropdown( | |
| choices=["Xception","InceptionV3","InceptionResNetV2","DenseNet201"], | |
| type="value", | |
| label="Select a model" | |
| ) | |
| ], | |
| outputs='label' | |
| ) | |
| # Launch the interface | |
| interface.launch() |