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(""" """) 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()