Spaces:
Sleeping
Sleeping
| import io | |
| import ast | |
| import base64 | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from PIL import Image | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from sklearn.preprocessing import normalize | |
| from transformers import CLIPModel, CLIPProcessor | |
| # ── App configuration ──────────────────────────────────────────────────────── | |
| MODEL_NAME = "openai/clip-vit-base-patch32" | |
| PARQUET_FILE = "nutrimatchemb.parquet" | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # ── Ingredient name mapping ────────────────────────────────────────────────── | |
| INGREDIENT_NAMES = { | |
| 0: 'background', 1: 'candy', 2: 'egg tart', | |
| 3: 'french fries', 4: 'chocolate', 5: 'biscuit', | |
| 6: 'popcorn', 7: 'pudding', 8: 'ice cream', | |
| 9: 'cheese/butter', 10: 'cake', 11: 'wine', | |
| 12: 'milkshake', 13: 'coffee', 14: 'juice', | |
| 15: 'milk', 16: 'tea', 17: 'almond', | |
| 18: 'red beans', 19: 'cashew', 20: 'dried cranberries', | |
| 21: 'soy', 22: 'walnut', 23: 'peanut', | |
| 24: 'egg', 25: 'apple', 26: 'date', | |
| 27: 'apricot', 28: 'avocado', 29: 'banana', | |
| 30: 'strawberry', 31: 'cherry', 32: 'blueberry', | |
| 33: 'raspberry', 34: 'mango', 35: 'olives', | |
| 36: 'peach', 37: 'lemon', 38: 'pear', | |
| 39: 'fig', 40: 'pineapple', 41: 'grape', | |
| 42: 'kiwi', 43: 'melon', 44: 'orange', | |
| 45: 'watermelon', 46: 'steak', 47: 'pork', | |
| 48: 'chicken/duck', 49: 'sausage', 50: 'fried meat', | |
| 51: 'lamb', 52: 'sauce', 53: 'crab', | |
| 54: 'fish', 55: 'shellfish', 56: 'shrimp', | |
| 57: 'soup', 58: 'bread', 59: 'corn', | |
| 60: 'burger', 61: 'pizza', 62: 'baozi/bun', | |
| 63: 'wonton/dumplings', 64: 'pasta', 65: 'noodles', | |
| 66: 'rice', 67: 'pie', 68: 'tofu', | |
| 69: 'eggplant', 70: 'potato', 71: 'garlic', | |
| 72: 'cauliflower', 73: 'tomato', 74: 'kelp', | |
| 75: 'seaweed', 76: 'spring onion', 77: 'rapeseed', | |
| 78: 'ginger', 79: 'okra', 80: 'lettuce', | |
| 81: 'pumpkin', 82: 'cucumber', 83: 'white radish', | |
| 84: 'carrot', 85: 'asparagus', 86: 'bamboo shoots', | |
| 87: 'broccoli', 88: 'celery', 89: 'cilantro/mint', | |
| 90: 'snow peas', 91: 'cabbage', 92: 'bean sprouts', | |
| 93: 'onion', 94: 'pepper', 95: 'green beans', | |
| 96: 'french beans', 97: 'king oyster mushroom', 98: 'shiitake', | |
| 99: 'enoki mushroom', 100: 'oyster mushroom', 101: 'button mushroom', | |
| 102: 'salad', 103: 'other ingredients', | |
| } | |
| # ── Load model and embedding database ──────────────────────────────────────── | |
| print("Loading CLIP model...") | |
| clip_model = CLIPModel.from_pretrained(MODEL_NAME).to(device) | |
| clip_processor = CLIPProcessor.from_pretrained(MODEL_NAME) | |
| clip_model.eval() | |
| print("Loading embedding database...") | |
| df_emb = pd.read_parquet(PARQUET_FILE) | |
| emb_matrix = np.vstack(df_emb["embedding"].values).astype("float32") | |
| emb_matrix = normalize(emb_matrix) | |
| print(f"Device: {device}") | |
| print(f"Embedding database shape: {emb_matrix.shape}") | |
| # ── Helper functions ───────────────────────────────────────────────────────── | |
| def safe_parse_ingredients(x): | |
| """Parse ingredient IDs whether stored as list, numpy array, or string.""" | |
| if isinstance(x, list): | |
| return x | |
| if isinstance(x, np.ndarray): | |
| return x.tolist() | |
| if isinstance(x, str): | |
| try: | |
| return ast.literal_eval(x) | |
| except Exception: | |
| return [] | |
| return [] | |
| def extract_features(output): | |
| """ | |
| Extract the actual embedding tensor from different possible CLIP outputs. | |
| This prevents the BaseModelOutputWithPooling .norm error. | |
| """ | |
| if isinstance(output, torch.Tensor): | |
| return output | |
| if hasattr(output, "image_embeds"): | |
| return output.image_embeds | |
| if hasattr(output, "pooler_output"): | |
| return output.pooler_output | |
| raise TypeError(f"Unsupported model output type: {type(output)}") | |
| def image_to_embedding(image): | |
| """Convert uploaded PIL image into a normalized CLIP image embedding.""" | |
| image = image.convert("RGB") | |
| inputs = clip_processor( | |
| images=image, | |
| return_tensors="pt" | |
| ).to(device) | |
| with torch.no_grad(): | |
| output = clip_model.get_image_features(**inputs) | |
| features = extract_features(output) | |
| features = features / features.norm(dim=-1, keepdim=True) | |
| query_vec = features.cpu().numpy().astype("float32") | |
| if query_vec.shape[1] != emb_matrix.shape[1]: | |
| raise ValueError( | |
| f"Embedding size mismatch: query has {query_vec.shape[1]} dimensions, " | |
| f"but database has {emb_matrix.shape[1]} dimensions." | |
| ) | |
| return query_vec | |
| def b64_to_image(image_b64): | |
| """Convert base64 image string back into PIL image.""" | |
| raw = base64.b64decode(image_b64) | |
| return Image.open(io.BytesIO(raw)).convert("RGB") | |
| def recommend(image, top_k=3): | |
| """Return top-k visually similar meals.""" | |
| if image is None: | |
| return [], "Please upload a food image." | |
| query_vec = image_to_embedding(image) | |
| scores = cosine_similarity(query_vec, emb_matrix)[0] | |
| top_indices = np.argsort(scores)[::-1][:top_k] | |
| gallery_items = [] | |
| text_lines = [] | |
| for rank, idx in enumerate(top_indices, start=1): | |
| row = df_emb.iloc[int(idx)] | |
| score = float(scores[idx]) | |
| ids = safe_parse_ingredients(row["ingredients"]) | |
| ingredient_names = [ | |
| INGREDIENT_NAMES.get(int(cid), "unknown") | |
| for cid in ids | |
| if int(cid) != 0 | |
| ] | |
| result_img = b64_to_image(row["image_b64"]) | |
| caption = ( | |
| f"Match #{rank} | Similarity: {score:.1%}\n" | |
| f"Ingredients: {', '.join(ingredient_names[:5])}" | |
| ) | |
| gallery_items.append((result_img, caption)) | |
| text_lines.append(caption) | |
| summary = "\n\n".join(text_lines) | |
| return gallery_items, summary | |
| # ── Gradio interface ───────────────────────────────────────────────────────── | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# NutriMatch AI") | |
| gr.Markdown( | |
| "Upload a meal image and get the top 3 visually similar meals. " | |
| "The system uses CLIP image embeddings and cosine similarity." | |
| ) | |
| gr.Markdown("## Presentation Video") | |
| gr.HTML(""" | |
| <iframe width="560" height="315" | |
| src="https://www.youtube.com/embed/9Fly6Usx5WU" | |
| title="NutriMatch AI Presentation Video" | |
| frameborder="0" | |
| allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" | |
| allowfullscreen> | |
| </iframe> | |
| """) | |
| gr.Markdown("## Food Recommendation App") | |
| input_image = gr.Image( | |
| type="pil", | |
| label="Upload a meal image" | |
| ) | |
| submit_button = gr.Button("Find Similar Meals") | |
| output_gallery = gr.Gallery( | |
| label="Top 3 visually similar meals", | |
| columns=3, | |
| height="auto" | |
| ) | |
| output_text = gr.Textbox( | |
| label="Recommendation details", | |
| lines=8 | |
| ) | |
| submit_button.click( | |
| fn=recommend, | |
| inputs=input_image, | |
| outputs=[output_gallery, output_text] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |