hjianganthony's picture
Create app.py
32f60c0
raw
history blame contribute delete
No virus
1.41 kB
import gradio as gr
from keras.models import load_model
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
from tensorflow.keras.utils import img_to_array
from tensorflow.keras.applications.resnet50 import preprocess_input
from PIL import Image
# model path
cnn_model = load_model('./model/cnn_model.h5')
resnet_model = load_model('./model/resnet_model.h5')
import json
with open('data/class_dict.json', 'r') as json_file:
class_dict = json.load(json_file)
# get class names
class_names = [class_dict[i] for i in sorted(class_dict.keys())]
def classify_insect(model_name, img):
img = img.resize((150, 150))
img_array = img_to_array(img)
img_array = np.expand_dims(img_array, axis=0)
img_preprocessed = preprocess_input(img_array)
# Select the model based on the dropdown choice
if model_name == "CNN Model":
model = cnn_model
elif model_name == "Transfer Learning ResNet":
model = resnet_model
# Make a prediction
prediction = model.predict(img_preprocessed)
return {class_name: float(score) for class_name, score in zip(class_names, prediction[0])}
iface = gr.Interface(
fn=classify_insect,
inputs=[
gr.Dropdown(choices=["CNN Model", "Transfer Learning ResNet"], label="Select Model"),
gr.Image(shape=(150,150))
],
outputs=gr.Label(num_top_classes=3)
)
iface.launch(share=True)