food-classification-101 / model_predictor.py
mhamza-007's picture
Added FastApi Backend
e3800b0
import tensorflow as tf
import numpy as np
model = tf.keras.models.load_model("model/101_food_class_fine_tuned_model.keras")
class_names = [
"apple_pie", "baby_back_ribs", "baklava", "beef_carpaccio", "beef_tartare", "beet_salad", "beignets", "bibimbap",
"bread_pudding", "breakfast_burrito", "bruschetta", "caesar_salad", "cannoli", "caprese_salad", "carrot_cake",
"ceviche", "cheesecake", "cheese_plate", "chicken_curry", "chicken_quesadilla", "chicken_wings", "chocolate_cake",
"chocolate_mousse", "churros", "clam_chowder", "club_sandwich", "crab_cakes", "creme_brulee", "croque_madame",
"cup_cakes", "deviled_eggs", "donuts", "dumplings", "edamame", "eggs_benedict", "escargots", "falafel",
"filet_mignon", "fish_and_chips", "foie_gras", "french_fries", "french_onion_soup", "french_toast",
"fried_calamari", "fried_rice", "frozen_yogurt", "garlic_bread", "gnocchi", "greek_salad", "grilled_cheese_sandwich",
"grilled_salmon", "guacamole", "gyoza", "hamburger", "hot_and_sour_soup", "hot_dog", "huevos_rancheros", "hummus",
"ice_cream", "lasagna", "lobster_bisque", "lobster_roll_sandwich", "macaroni_and_cheese", "macarons", "miso_soup",
"mussels", "nachos", "omelette", "onion_rings", "oysters", "pad_thai", "paella", "pancakes", "panna_cotta",
"peking_duck", "pho", "pizza", "pork_chop", "poutine", "prime_rib", "pulled_pork_sandwich", "ramen", "ravioli",
"red_velvet_cake", "risotto", "samosa", "sashimi", "scallops", "seaweed_salad", "shrimp_and_grits",
"spaghetti_bolognese", "spaghetti_carbonara", "spring_rolls", "steak", "strawberry_shortcake", "sushi", "tacos",
"takoyaki", "tiramisu", "tuna_tartare", "waffles"
]
def prediction(img):
img = tf.convert_to_tensor(img, dtype=tf.float32)
img = tf.expand_dims(img, axis=0)
pred = model.predict(img)
# Get the top 10 predictions
top_indices = np.argsort(pred[0])[::-1][:10] # Sort probabilities in descending order
top_probs = pred[0][top_indices] # Get top probabilities
top_classes = [class_names[i] for i in top_indices] # Get corresponding class names
predictions_list = [{"class": top_classes[i], "probability": float(top_probs[i])} for i in range(len(top_classes))]
return predictions_list