Spaces:
Running
Running
| import streamlit as st | |
| import streamlit.components.v1 as components | |
| import pickle | |
| import sentence_transformers | |
| from sentence_transformers import SentenceTransformer, util | |
| from PIL import Image | |
| import torch | |
| import spacy | |
| import os | |
| import glob | |
| import random | |
| torch.set_num_threads(4) | |
| def get_spacy_dbpedia_highlights(ingredients): | |
| import spacy | |
| import spacy_dbpedia_spotlight | |
| raw_ingredients = ingredients | |
| import re | |
| ingredients = re.sub("[0-9,()\/\-\.]", "", ingredients) | |
| doc = nlp(ingredients) | |
| for ent in doc.ents: | |
| if ent.text.lower() not in stop_words and ent.text in raw_ingredients: | |
| replace_str = '<mark style="color: green; background-color:yellow"> <a href="' + ent.kb_id_ + '" target="_blank"> ' + ent.text + '</a> </mark>' | |
| raw_ingredients = raw_ingredients.replace(ent.text, replace_str) | |
| return raw_ingredients | |
| def detect_food(query, text_emb, labels, k=1): | |
| query_emb = model.encode(Image.open(query), convert_to_tensor=True, show_progress_bar=False) | |
| hits = util.semantic_search(query_emb, text_emb, top_k=k)[0] | |
| results = [] | |
| for i, hit in enumerate(hits): | |
| results.append((labels[hit['corpus_id']], hit['score'])) | |
| if i > 2: | |
| break | |
| return results | |
| def run_search(food_image, col2): | |
| with open("./Pretrained/labels.pkl", 'rb') as fIn: | |
| labels = pickle.load(fIn) | |
| emb_filename = './Pretrained/food_embeddings.pkl' | |
| text_emb = torch.load(emb_filename, map_location=torch.device('cpu')) | |
| results = detect_food(food_image, text_emb, labels, 3) | |
| food_recognised, score = results[0] | |
| del text_emb | |
| del labels | |
| import pysos | |
| id2recipe = pysos.Dict("./Pretrained/id2recipe") | |
| food2id = pysos.Dict("./Pretrained/food2id") | |
| id = food2id[food_recognised] | |
| recipe_name = food_recognised.title() | |
| ingredients_list =id2recipe[id]['ingredients'] | |
| highlighted_ingredients= get_spacy_dbpedia_highlights(ingredients_list) | |
| recipe= id2recipe[id]['instructions'] | |
| dataset = " " + id2recipe[id]['dataset'] | |
| if dataset.strip() == "Recipe1M": | |
| nutritional_facts= "For nutritional facts, schedule and servings, visit the link in the footer" | |
| else: | |
| nutritional_facts = id2recipe[id]['nutrition_facts'] | |
| source= id2recipe[id]['recipesource'] | |
| del id2recipe | |
| del food2id | |
| st.markdown("<br/>", unsafe_allow_html=True) | |
| with col2: | |
| st.markdown("<b>Top 3 predictions   </b>", unsafe_allow_html=True) | |
| results_static_tag = '<html><title>W3.CSS</title><meta name="viewport" content="width=device-width, initial-scale=1"><link rel="stylesheet" href="https://www.w3schools.com/w3css/4/w3.css"><body><div class="w3-container">{}</div></body></html>' | |
| result_rows = "" | |
| for i, result in enumerate(results): | |
| results_dynamic_tag= '{} <br/> <div class="w3-light-grey"> <div class="{}" style="height:4px;width:{}%"></div> </div><br>' | |
| if i == 0: | |
| results_dynamic_tag = results_dynamic_tag.format("<b>" + str(i+1) + "." + result[0].title() + "</b>", 'w3-blue', result[1] * 100) | |
| else: | |
| results_dynamic_tag = results_dynamic_tag.format(str(i+1) + "." + result[0].title(), "w3-orange" ,result[1] * 100) | |
| result_rows += results_dynamic_tag | |
| results_static_tag = results_static_tag.format(result_rows) | |
| st.markdown(results_static_tag, unsafe_allow_html=True) | |
| title_tag = '<h4> Recipe for top result:  ' + recipe_name + '</h4>' | |
| st.markdown(title_tag, unsafe_allow_html=True) | |
| ing_hdr_tag = '<h5> Ingredients </h5>' | |
| ing_style= "{border: 3x outset white; background-color: #ccf5ff; color: black; text-align: left; font-size: 14px; padding: 5px;}" | |
| ing_tag = '<html><head><style>.ingdiv{}</style></head><body><div class="ingdiv">{}</div></body></html>' | |
| ing_tag = ing_tag.format(ing_style, highlighted_ingredients.strip()) | |
| st.markdown(ing_hdr_tag, unsafe_allow_html=True) | |
| st.markdown(ing_tag + "<br/>", unsafe_allow_html=True) | |
| rec_hdr_tag = '<h5> Recipe </h5>' | |
| rec_style= "{border: 3x outset white; background-color: #ffeee6; color: black; text-align: left; font-size: 14px; padding: 5px;}" | |
| rec_tag = '<html><head><style>.recdiv{}</style></head><body><div class="recdiv">{}</div></body></html>' | |
| rec_tag = rec_tag.format(rec_style, recipe.strip()) | |
| st.markdown(rec_hdr_tag, unsafe_allow_html=True) | |
| st.markdown(rec_tag + "<br/>", unsafe_allow_html=True) | |
| src_hdr_tag = '<h5> Recipe source </h5>' | |
| src_tag = '<a href={} target="_blank">{}</a>' | |
| src_tag = src_tag.format(source, source) | |
| st.markdown(src_hdr_tag, unsafe_allow_html=True) | |
| st.markdown(src_tag + "<br/>", unsafe_allow_html=True) | |
| return 1 | |
| if 'models_loaded' not in st.session_state: | |
| st.session_state['models_loaded'] = False | |
| st.title('WTF - What The Food 🤬') | |
| st.subheader("Image to Recipe - 1.5M foods supported") | |
| st.markdown("Built for fun with 💙 by a quintessential foodie - Prithivi Da | [@prithivida](https://twitter.com/prithivida) |[[GitHub]](https://github.com/PrithivirajDamodaran) <br/> <hr style='height:1px;border:none;color:violet;background-color:gray;' />", unsafe_allow_html=True) | |
| st.write("""Read Me: The goal is to detect a "Single food item" from the image and retrieve it's recipe. So by design the model works well on single foods. It works on platters too fx English breakfast but it may not perform well on a custom combination with multiple recipes or hyper-local foods. | |
| """) | |
| def load_image(image_file): | |
| img = Image.open(image_file) | |
| return img | |
| def load_models(): | |
| with st.spinner(text="Loading Models..."): | |
| os.system("python -m spacy download en_core_web_sm") | |
| nlp = spacy.load('en_core_web_sm') | |
| nlp.add_pipe('dbpedia_spotlight') | |
| model = SentenceTransformer('clip-ViT-B-32') | |
| stop_words = set(['chopped', 'freshly ground', 'freshly squeezed', 'dash', 'powder', 'rice', 'ice', 'noodles', 'pepper', 'milk', 'ced', 'cheese', 'sugar', 'salt', 'pkt', 'minced', 'onion', 'onions', 'garlic', 'butter', 'slices', 'ounce', 'sauce', 'freshly', 'grated', 'teaspoon', 'cup', 'oz', '⁄', 'to', 'or', 'diced', 'into', 'pound', 'dried', 'water', 'about', 'whole', 'small', 'vegetable', 'inch', 'tbsp', 'cooked', 'large', 'sliced', 'dry', 'optional', 'package', 'ounces', 'unsalted', 'lbs', 'green', 'flour', 'for', 'wine', 'crushed', 'drained', 'lb', 'frozen', 'tsp', 'finely', 'medium', 'tablespoon', 'tablespoons', 'juice', 'shredded', 'can', 'minced', 'fresh', 'cut', 'pieces', 'in', 'thinly', 'of', 'extract', 'teaspoons', 'ground', 'and', 'cups', 'peeled', 'taste', 'ml', 'lengths']) | |
| st.session_state['nlp'] = nlp | |
| st.session_state['model'] = model | |
| st.session_state['stop_words'] = stop_words | |
| if not st.session_state['models_loaded']: | |
| load_models() | |
| st.session_state['models_loaded'] = True | |
| random_button = st.button('⚡ Try a Random Food') | |
| st.write("(or)") | |
| image_file = st.file_uploader("Tip: Upload HD images for better results.", type=["jpg","jpeg"]) | |
| nlp = st.session_state['nlp'] | |
| model = st.session_state['model'] | |
| stop_words = st.session_state['stop_words'] | |
| col1, col2 = st.columns(2) | |
| if random_button: | |
| with st.spinner(text="Detecting food..."): | |
| samples = glob.glob('./samples' + "/*") | |
| random_sample = random.choice(samples) | |
| pil_image = load_image(random_sample) | |
| with col1: | |
| st.image(pil_image, use_column_width='auto') | |
| return_code = run_search(random_sample, col2) | |
| else: | |
| if image_file is not None: | |
| pil_image = load_image(image_file) | |
| with open(image_file.name, 'wb') as f: | |
| pil_image.save(f) | |
| with col1: | |
| st.image(pil_image, use_column_width='auto') | |
| with st.spinner(text="Detecting food..."): | |
| return_code = run_search(image_file.name, col2) | |
| os.system('rm -r "' + image_file.name + '"') | |