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