Spaces:
Sleeping
Sleeping
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
from PIL import Image | |
# Load your trained models | |
model1 = tf.keras.models.load_model('model/FoodVisionFineTuneAug/') | |
model2 = tf.keras.models.load_model('model/FoodVisionFineTune/') | |
with open('classes.txt', 'r') as f: | |
classes = [line.strip() for line in f] | |
# Add information about the models | |
model1_info = """ | |
### Model 1 Information | |
This model is based on the EfficientNetB0 architecture and was trained on the Food101 dataset. | |
""" | |
model2_info = """ | |
### Model 2 Information | |
This model is based on the EfficientNetB0 architecture and was trained on augmented data, providing improved generalization. | |
""" | |
def preprocess(image: Image.Image): | |
# Convert numpy array to PIL Image | |
image = Image.fromarray((image * 255).astype(np.uint8)) | |
image = image.resize((224, 224)) # replace with the input size of your models | |
image = np.array(image) | |
# image = image / 255.0 # normalize if you've done so while training | |
image = np.expand_dims(image, axis=0) | |
return image | |
def predict(model_selection, image: Image.Image): | |
# Choose the model based on the dropdown selection | |
model = model1 if model_selection == "EfficentNetB0 Fine Tune" else model2 | |
image = preprocess(image) | |
prediction = model.predict(image) | |
predicted_class = classes[np.argmax(prediction)] | |
confidence = np.max(prediction) | |
return predicted_class, confidence | |
iface = gr.Interface( | |
fn=predict, | |
inputs=[gr.Dropdown(["EfficentNetB0 Fine Tune", "EfficentNetB0 Fine Tune Augmented"]), gr.Image()], | |
outputs=[gr.Textbox(label="Predicted Class"), gr.Textbox(label="Confidence")], | |
title="Transfer Learning Mini Project", | |
description=f"{model1_info}\n\n{model2_info}", | |
) | |
iface.launch() | |