File size: 6,531 Bytes
43b5226 5ddf592 43b5226 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import torch
from torch import nn
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import time
from typing import Tuple, Dict
import json
import gradio as gr
# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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', 'cheese_plate', 'cheesecake', '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'
]
# Function to load JSON files and concatenate them
def load_and_concatenate_json(json_paths: Tuple[str, str]) -> list:
combined_data = []
for json_path in json_paths:
with open(json_path, 'r') as f:
data = json.load(f)
combined_data.extend(data)
return combined_data
json_paths = ('./demos/recipenet/recipesData_v1.json', './demos/recipenet/recipesData_v2.json') # Replace with your JSON file paths
combined_data = load_and_concatenate_json(json_paths)
# Load the model
weights = models.EfficientNet_B0_Weights.DEFAULT
model = models.efficientnet_b0(weights=weights)
# Update the classifier and heads
model.classifier = nn.Sequential(
nn.Dropout(p=0.2, inplace=True),
nn.Linear(in_features=1280, out_features=101)
)
model.heads = nn.Linear(in_features=768, out_features=101)
# Load the state dictionary
state_dict = torch.load('./demos/recipenet/best_model.pth', map_location=device)
model.load_state_dict(state_dict)
model.to(device)
auto_transforms = weights.transforms()
def get_recipe_details(label: str, combined_data: list) -> dict:
label = label.replace("_", " ")
for item in combined_data:
if item["FoodName"].lower() == label.lower():
return {
"Ingredients": item["RawIngredients"],
"RecipeURL": item["RecipeURL"],
"RecipeName": item["RecipeName"]
}
return {}
def predict(img: Image.Image, model, class_names, device, combined_data: list) -> Tuple[Dict[str, float], float, dict]:
# Start a timer
start_time = time.time()
# Transform the input image for use with EfficientNet_B0
img = auto_transforms(img).unsqueeze(0).to(device)
# Put the model in evaluation mode
model.eval()
with torch.inference_mode():
# Make predictions
pred_probs = torch.softmax(model(img), dim=1)
# Create a dictionary of prediction labels and probabilities
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
# End the timer
end_time = time.time()
pred_time = end_time - start_time
# Get the maximum prediction label and probability
max_label = max(pred_labels_and_probs, key=pred_labels_and_probs.get)
max_prob = pred_labels_and_probs[max_label]
# Fetch recipe details for the predicted label
recipe_details = get_recipe_details(max_label, combined_data)
return pred_labels_and_probs, pred_time, recipe_details
def predict_wrapper(img: Image.Image):
# Provide the additional arguments required by the predict function
pred_labels_and_probs, pred_time, recipe_details = predict(img, model, class_names, device, combined_data)
# Extract the recipe, website, and ingredients from recipe_details
recipe = recipe_details.get('RecipeName', '')
website = recipe_details.get('RecipeURL', '')
ingredients = recipe_details.get('Ingredients', [])
time = len(ingredients) * 10 # Assuming each ingredient takes 10 seconds to prepare
ingredients = ', '.join(ingredients) # Convert the list of ingredients to a string
return pred_labels_and_probs, pred_time, recipe, time, website, ingredients
# Define the CSS for transparent background
css = """
.gradio-container,.svelte-vt1mxs.gap.panel {
background: repeating-linear-gradient(
to top,
rgba(255, 255, 255, 0.03) 0px 2px,
transparent 2px 4px
),
linear-gradient(to bottom, #200933 75%, #3d0b43);
color: #d5d7de;
font-family: sans-serif;
}
gradio-app[control_page_title="true"][embed="false"][eager="true"] {
background-color: #271139 !important;
background-image: linear-gradient(to bottom, transparent 95%, #000 5%);
}
.gradio-root {
overflow: hidden;
box-shadow: 0px 0px 10px 0px rgba(0, 0, 0, 0.1);
}
.gradio-control {
border: none;
}
.gradio-input {
background-color: rgba(255, 255, 255, 0.1);
border: 1px solid rgba(255, 255, 255, 0.2);
}
"""
demo = gr.Interface(
fn=predict_wrapper,
inputs=gr.Image(type='pil'),
outputs=[
gr.Label(num_top_classes=3, label='Predictions'),
gr.Number(label='Prediction Time (s)'),
gr.Textbox(label='Dish Name', elem_id="textbox_id"),
gr.Textbox(label='Preparation time(in Minutes)'),
gr.Textbox(label='Recipe available at'),
gr.Textbox(label='Ingredients')
],
description="recipenet",
article="Made using EfficientNet-b0",
css=css, # Add the custom CSS
examples=[["./demos/recipenet/examples/pic1.jpg"], ["./demos/recipenet/examples/pic2.jpg"], ["./demos/recipenet/examples/pic3.jpg"]]
)
demo.launch(debug=False, share=False)
|