Update src/streamlit_app.py

src/streamlit_app.py

@@ -10,7 +10,11 @@ from transformers import (
 )
 
 # Set Streamlit UI
-st.set_page_config(page_title="🍽️ Food Nutrition Estimator", page_icon="🥗", layout="centered")
+st.set_page_config(
+    page_title="🍽️ Food Nutrition Estimator",
+    page_icon="🥗",
+    layout="centered"
+)
 st.title("🍽️ Food Nutrition Estimator")
 st.markdown("Upload a food image and get nutritional information generated by AI!")
 
@@ -25,7 +29,8 @@ manual_transform = transforms.Compose([
     transforms.Resize(224),
     transforms.CenterCrop(196),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
     transforms.ConvertImageDtype(torch.float32)
 ])
 
@@ -64,4 +69,158 @@ def load_models():
 
 model_convnext, tokenizer, model_llm, device = load_models()
 
-# ... rest of your code remains unchanged ...
+# 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
+uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    try:
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="Uploaded Image", use_column_width=True)
+
+        # Predict with ConvNeXt
+        input_tensor = manual_transform(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            outputs = model_convnext(pixel_values=input_tensor)
+        pred_idx = outputs.logits.argmax(-1).item()
+        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]}"
+        )
+        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.'"
+        )
+        st.subheader("🧾 Nutrition Information")
+        st.write(f"🤖 Prompt: Describe {pred_label} nutrition in your own words based on provided data.")
+
+        input_ids = tokenizer(prompt, return_tensors="pt").to(model_llm.device)
+        with torch.no_grad():
+            output = model_llm.generate(
+                **input_ids,
+                max_new_tokens=300,
+                temperature=0.8,
+                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
+
+        st.info(caption)
+
+    except Exception as e:
+        st.error(f"Something went wrong: {e}")
+
+# Footer
+st.markdown("---")
+st.markdown("Built with ❤️ using Streamlit and Hugging Face by **shingguy1**")