Update src/streamlit_app.py

src/streamlit_app.py

@@ -46,7 +46,7 @@ st.sidebar.markdown("""
 def load_models():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # ConvNeXt for classification (updated repo)
+    # ConvNeXt for classification
     model_convnext = ConvNextForImageClassification.from_pretrained(
         "shingguy1/fine_tuned_convnext",
         cache_dir=cache_dir,
@@ -69,99 +69,16 @@ def load_models():
 
 model_convnext, tokenizer, model_llm, device = load_models()
 
-# Fallback nutritional data for all ten ConvNeXt food classes
-fallback_nutrition = {
-    "hamburger": (
-        "A standard hamburger (single beef patty, bun, lettuce, tomato, no condiments) weighs about 150g per serving. "
-        "It has approximately 300 calories, 20g protein, 30g carbohydrates, and 12g fat. "
-        "Main ingredients: ground beef patty (80/20 lean-to-fat ratio), white bun, lettuce, tomato. "
-        "Cooking method: grilled or pan-fried. Nutritional facts: 500mg sodium, 10% daily value iron, minimal fiber. "
-        "Add-ins like cheese (+100 calories, 7g fat) or mayo (+90 calories, 10g fat) increase calories and fat. "
-        "Substitute chicken patty for lower fat (8g)."
-    ),
-    "pizza": (
-        "A standard cheese pizza slice (1/8 of a 14-inch pizza) weighs about 100g per serving. "
-        "It has approximately 270 calories, 12g protein, 34g carbohydrates, and 10g fat. "
-        "Main ingredients: dough, tomato sauce, mozzarella cheese. "
-        "Cooking method: baked in an oven. Nutritional facts: 600mg sodium, 15% daily value calcium, 2g fiber. "
-        "Add-ins like pepperoni (+50 calories, 5g fat) or extra cheese (+80 calories, 6g fat) increase calories. "
-        "Substitute cauliflower crust for lower carbs (20g)."
-    ),
-    "sushi": (
-        "A standard sushi roll (6 pieces, e.g., California roll) weighs about 150g per serving. "
-        "It has approximately 200 calories, 7g protein, 30g carbohydrates, and 5g fat. "
-        "Main ingredients: sushi rice, nori (seaweed), crab or imitation crab, avocado, cucumber. "
-        "Cooking method: raw or assembled without cooking. Nutritional facts: 400mg sodium, 10% daily value vitamin A. "
-        "Add-ins like soy sauce (+200mg sodium) or spicy mayo (+100 calories, 10g fat) increase sodium or fat. "
-        "Substitute brown rice for higher fiber (2g)."
-    ),
-    "ceasar_salad": (
-        "A standard garden salad (mixed greens, tomato, cucumber, no dressing) weighs about 200g per serving. "
-        "It has approximately 50 calories, 2g protein, 10g carbohydrates, and 0.5g fat. "
-        "Main ingredients: lettuce, spinach, tomato, cucumber, carrots. "
-        "Cooking method: raw, no cooking. Nutritional facts: 50mg sodium, 50% daily value vitamin C, 4g fiber. "
-        "Add-ins like ranch dressing (+150 calories, 15g fat) or grilled chicken (+120 calories, 20g protein) increase calories or protein. "
-        "Substitute vinaigrette for lower fat (5g)."
-    ),
-    "pasta": (
-        "A standard pasta dish (1 cup cooked spaghetti with marinara sauce) weighs about 200g per serving. "
-        "It has approximately 220 calories, 7g protein, 43g carbohydrates, and 2g fat. "
-        "Main ingredients: wheat pasta, tomato sauce, olive oil. "
-        "Cooking method: boiled pasta, sauce simmered. Nutritional facts: 400mg sodium, 10% daily value vitamin A, 3g fiber. "
-        "Add-ins like meatballs (+200 calories, 15g fat) or parmesan (+50 calories, 4g fat) increase calories. "
-        "Substitute whole-grain pasta for higher fiber (5g)."
-    ),
-    "ice_cream": (
-        "A standard ice cream serving (1/2 cup vanilla) weighs about 100g. "
-        "It has approximately 200 calories, 4g protein, 20g carbohydrates, and 12g fat. "
-        "Main ingredients: cream, sugar, milk, vanilla extract. "
-        "Cooking method: churned and frozen. Nutritional facts: 100mg sodium, 15% daily value calcium, 0g fiber. "
-        "Add-ins like chocolate syrup (+100 calories, 2g fat) or sprinkles (+50 calories) increase calories. "
-        "Substitute frozen yogurt for lower fat (7g)."
-    ),
-    "fried_rice": (
-        "A standard fried rice serving (1 cup with vegetables and egg) weighs about 200g. "
-        "It has approximately 250 calories, 8g protein, 35g carbohydrates, and 9g fat. "
-        "Main ingredients: white rice, egg, peas, carrots, soy sauce, vegetable oil. "
-        "Cooking method: stir-fried. Nutritional facts: 700mg sodium, 10% daily value vitamin A, 2g fiber. "
-        "Add-ins like chicken (+100 calories, 15g protein) or shrimp (+80 calories, 12g protein) increase protein. "
-        "Substitute brown rice for higher fiber (3g)."
-    ),
-    "tacos": (
-        "A standard taco (1 soft corn tortilla with beef, lettuce, cheese) weighs about 100g per serving. "
-        "It has approximately 200 calories, 10g protein, 15g carbohydrates, and 10g fat. "
-        "Main ingredients: ground beef, corn tortilla, lettuce, cheddar cheese, salsa. "
-        "Cooking method: beef pan-fried, tortilla warmed. Nutritional facts: 400mg sodium, 10% daily value calcium, 2g fiber. "
-        "Add-ins like sour cream (+50 calories, 5g fat) or guacamole (+80 calories, 7g fat) increase fat. "
-        "Substitute fish for lower fat (6g)."
-    ),
-    "steak": (
-        "A standard steak (4 oz grilled sirloin) weighs about 113g per serving. "
-        "It has approximately 250 calories, 25g protein, 0g carbohydrates, and 15g fat. "
-        "Main ingredients: beef sirloin, salt, pepper. "
-        "Cooking method: grilled or pan-seared. Nutritional facts: 300mg sodium, 20% daily value iron, 0g fiber. "
-        "Add-ins like butter sauce (+100 calories, 10g fat) or mashed potatoes (+150 calories, 5g fat) increase calories. "
-        "Substitute leaner cut (e.g., filet mignon) for lower fat (10g)."
-    ),
-    "chocolate_cake": (
-        "A standard chocolate cake slice (1/12 of a 9-inch cake) weighs about 100g per serving. "
-        "It has approximately 350 calories, 5g protein, 50g carbohydrates, and 15g fat. "
-        "Main ingredients: flour, sugar, cocoa powder, butter, eggs. "
-        "Cooking method: baked. Nutritional facts: 200mg sodium, 5% daily value iron, 2g fiber. "
-        "Add-ins like frosting (+100 calories, 5g fat) or whipped cream (+50 calories, 5g fat) increase calories. "
-        "Substitute gluten-free flour for dietary needs (same calories)."
-    )
-}
-
-# Upload image
+# Image uploader
 uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
     try:
+        # Load and display
         image = Image.open(uploaded_file).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
 
-        # Predict with ConvNeXt
+        # Predict food label
         input_tensor = manual_transform(image).unsqueeze(0).to(device)
         with torch.no_grad():
             outputs = model_convnext(pixel_values=input_tensor)
@@ -169,32 +86,14 @@ if uploaded_file is not None:
         pred_label = model_convnext.config.id2label[pred_idx]
         st.success(f"🍴 Predicted Food: **{pred_label}**")
 
-        # Generate nutrition caption using TinyLlama
-        fallback_text = fallback_nutrition.get(pred_label.lower(), "Nutritional facts unavailable for this item.")
-        simplified_fallback = (
-            f"Food: {pred_label}\n"
-            f"Serving size: {fallback_text.split('weighs about ')[1].split(' per serving')[0]}\n"
-            f"Calories: {fallback_text.split('approximately ')[1].split(' calories')[0]} calories\n"
-            f"Protein: {fallback_text.split('protein, ')[1].split('g ')[0]}g\n"
-            f"Carbs: {fallback_text.split('carbohydrates, ')[1].split('g ')[0]}g\n"
-            f"Fat: {fallback_text.split('and ')[1].split('g fat')[0]}g\n"
-            f"Ingredients: {fallback_text.split('Main ingredients: ')[1].split('. ')[0]}\n"
-            f"Cooking method: {fallback_text.split('Cooking method: ')[1].split('. ')[0]}\n"
-            f"Extras: {fallback_text.split('Add-ins like ')[1].split('. Substitute')[0]}\n"
-            f"Substitution: {fallback_text.split('Substitute ')[1].split('.')[0]}"
-        )
+        # Generate nutrition description with LLM
         prompt = (
-            f"Below is nutritional information for a {pred_label}:\n"
-            f"{simplified_fallback}\n\n"
-            f"Using this data, write a concise, natural description of the {pred_label}'s nutrition in your own words. "
-            f"Do not copy the text above; create a new description with different wording. "
-            f"Include calories, macronutrients (protein, carbs, fat), serving size, ingredients, cooking method, add-ins, and substitution. "
-            f"Example for a taco: 'A classic taco, around 100g, offers 200 calories, 10g protein, 15g carbs, and 10g fat. "
-            f"It’s prepared with ground beef, corn tortilla, lettuce, cheese, and salsa, with the beef fried and tortilla warmed. "
-            f"Add sour cream for 50 extra calories or switch to fish for 6g fat.'"
+            f"Please provide a concise nutritional overview for a {pred_label}. "
+            "Include typical serving size, approximate calories, macronutrient breakdown "
+            "(protein, carbs, fat), main ingredients, common cooking method, and one substitution suggestion."
         )
         st.subheader("🧾 Nutrition Information")
-        st.write(f"🤖 Prompt: Describe {pred_label} nutrition in your own words based on provided data.")
+        st.write(f"🤖 Prompt to LLM:\n\n{prompt}")
 
         input_ids = tokenizer(prompt, return_tensors="pt").to(model_llm.device)
         with torch.no_grad():
@@ -205,17 +104,7 @@ if uploaded_file is not None:
                 top_p=0.9,
                 do_sample=True
             )
-        input_len = input_ids.input_ids.shape[1]
-        caption = tokenizer.decode(output[0][input_len:], skip_special_tokens=True).strip()
-
-        # Check similarity to fallback
-        fallback_words = set(fallback_text.lower().split())
-        caption_words = set(caption.lower().split())
-        similarity = len(fallback_words & caption_words) / max(len(fallback_words), 1)
-
-        if not caption or len(caption) < 60 or any(kw in caption.lower() for kw in ["rephrase", "below is", "example for"]) or similarity > 0.8:
-            caption = fallback_text
-
+        caption = tokenizer.decode(output[0], skip_special_tokens=True).strip()
         st.info(caption)
 
     except Exception as e: