Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,8 +5,9 @@ st.set_page_config(
     layout="centered"
 )
 
-import torch
 import os
+import torch
+import random
 from PIL import Image
 import torchvision.transforms as transforms
 from transformers import (
@@ -22,108 +23,200 @@ def main():
     os.makedirs(cache_dir, exist_ok=True)
     os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
 
-    # 2. Image transform for ViT
+    # 2. Nutritional lookup table
+    nutritional_info = {
+        "pizza": {
+            "serving": "100 g (1 slice)",
+            "calories": "270 kcal",
+            "protein": "12 g",
+            "carbs": "34 g",
+            "fat": "10 g",
+            "ingredients": "dough, tomato sauce, mozzarella cheese",
+            "method": "baked",
+            "substitute": "cauliflower crust"
+        },
+        "hamburger": {
+            "serving": "150 g",
+            "calories": "300 kcal",
+            "protein": "20 g",
+            "carbs": "30 g",
+            "fat": "12 g",
+            "ingredients": "ground beef patty (80/20), bun, lettuce, tomato",
+            "method": "grilled or pan-fried",
+            "substitute": "chicken patty"
+        },
+        "sushi": {
+            "serving": "150 g (6 pieces)",
+            "calories": "200 kcal",
+            "protein": "7 g",
+            "carbs": "30 g",
+            "fat": "5 g",
+            "ingredients": "sushi rice, nori, crab (or imitation), avocado, cucumber",
+            "method": "assembled raw",
+            "substitute": "brown rice"
+        },
+        "salad": {
+            "serving": "200 g",
+            "calories": "50 kcal",
+            "protein": "2 g",
+            "carbs": "10 g",
+            "fat": "0.5 g",
+            "ingredients": "mixed greens, tomato, cucumber, carrots",
+            "method": "raw",
+            "substitute": "vinaigrette instead of ranch"
+        },
+        "pasta": {
+            "serving": "200 g (1 cup)",
+            "calories": "220 kcal",
+            "protein": "7 g",
+            "carbs": "43 g",
+            "fat": "2 g",
+            "ingredients": "wheat pasta, marinara sauce, olive oil",
+            "method": "boiled and simmered",
+            "substitute": "whole-grain pasta"
+        },
+        "ice_cream": {
+            "serving": "100 g (½ cup)",
+            "calories": "200 kcal",
+            "protein": "4 g",
+            "carbs": "20 g",
+            "fat": "12 g",
+            "ingredients": "cream, sugar, milk, vanilla",
+            "method": "churned and frozen",
+            "substitute": "frozen yogurt"
+        },
+        "fried_rice": {
+            "serving": "200 g (1 cup)",
+            "calories": "250 kcal",
+            "protein": "8 g",
+            "carbs": "35 g",
+            "fat": "9 g",
+            "ingredients": "rice, egg, peas, carrots, soy sauce, oil",
+            "method": "stir-fried",
+            "substitute": "brown rice"
+        },
+        "tacos": {
+            "serving": "100 g (1 soft taco)",
+            "calories": "200 kcal",
+            "protein": "10 g",
+            "carbs": "15 g",
+            "fat": "10 g",
+            "ingredients": "ground beef, corn tortilla, lettuce, cheese, salsa",
+            "method": "beef pan-fried, tortilla warmed",
+            "substitute": "fish filling"
+        },
+        "steak": {
+            "serving": "113 g (4 oz)",
+            "calories": "250 kcal",
+            "protein": "25 g",
+            "carbs": "0 g",
+            "fat": "15 g",
+            "ingredients": "beef sirloin, salt, pepper",
+            "method": "grilled or pan-seared",
+            "substitute": "leaner cut (filet mignon)"
+        },
+        "chocolate_cake": {
+            "serving": "100 g (1 slice)",
+            "calories": "350 kcal",
+            "protein": "5 g",
+            "carbs": "50 g",
+            "fat": "15 g",
+            "ingredients": "flour, sugar, cocoa, butter, eggs",
+            "method": "baked",
+            "substitute": "gluten-free flour"
+        }
+    }
+
+    # 3. Image transform for ViT
     manual_transform = transforms.Compose([
         transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.Lambda(lambda img: img.convert("RGB")),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                             std=[0.229, 0.224, 0.225]),
-        transforms.ConvertImageDtype(torch.float32)
+                             std=[0.229, 0.224, 0.225])
     ])
 
-    # 3. Sidebar info
+    # 4. Sidebar info
     st.sidebar.header("Models Used")
     st.sidebar.markdown("""
-    - 🖼️ **Image Classifier**: `shingguy1/fine_tuned_vit`  
-    - 💬 **Text Generator**: `google/flan-t5-small`
+    - 🖼️ **Image Classifier**: `shingguy1/fine_tuned_vit`
+    - 💬 **Paraphraser**: `google/flan-t5-small`
     """)
 
-    # 4. Load models (cached)
+    # 5. Load models (cached)
     @st.cache_resource
     def load_models():
-        device = torch.device("cpu")  # CPU-only environment
-
-        # ViT classifier
-        model_vit = ViTForImageClassification.from_pretrained(
+        device = torch.device("cpu")
+        vit = ViTForImageClassification.from_pretrained(
             "shingguy1/fine_tuned_vit",
             cache_dir=cache_dir,
             use_auth_token=hf_token
         ).to(device)
-
-        # FLAN-T5 Small for generation
-        tokenizer_llm = AutoTokenizer.from_pretrained(
+        tok = AutoTokenizer.from_pretrained(
             "google/flan-t5-small",
             cache_dir=cache_dir,
             use_auth_token=hf_token
         )
-        model_llm = T5ForConditionalGeneration.from_pretrained(
+        paraphraser = T5ForConditionalGeneration.from_pretrained(
             "google/flan-t5-small",
             cache_dir=cache_dir,
             use_auth_token=hf_token
         ).to(device)
+        return vit, tok, paraphraser, device
+
+    model_vit, tokenizer_t5, model_t5, device = load_models()
+
+    # 6. Uploader
+    uploaded = st.file_uploader("Upload a food image...", type=["jpg","png","jpeg"])
+    if uploaded:
+        img = Image.open(uploaded)
+        st.image(img, caption="Your Food", use_column_width=True)
+
+        # classify
+        inp = manual_transform(img).unsqueeze(0).to(device)
+        with torch.no_grad():
+            out = model_vit(pixel_values=inp)
+        label = model_vit.config.id2label[out.logits.argmax(-1).item()]
+        st.success(f"🍽️ Detected: **{label}**")
+
+        # lookup
+        data = nutritional_info.get(label.lower())
+        if not data:
+            st.error("No nutrition data for this item.")
+            return
+
+        # slot-fill template
+        templates = [
+            "A typical {label} serving ({serving}) contains about {calories}, with {protein} protein, {carbs} carbs, and {fat} fat. "
+            "Made from {ingredients} and usually {method}. Try {substitute} as a healthier swap.",
+            "For {label}, one {serving} provides {calories}. It offers {protein} protein, {carbs} carbohydrates, and {fat} fat. "
+            "Ingredients include {ingredients}, and it's {method}. You can substitute {substitute}."
+        ]
+        raw = random.choice(templates).format(label=label,
+                                              serving=data["serving"],
+                                              calories=data["calories"],
+                                              protein=data["protein"],
+                                              carbs=data["carbs"],
+                                              fat=data["fat"],
+                                              ingredients=data["ingredients"],
+                                              method=data["method"],
+                                              substitute=data["substitute"])
+
+        # paraphrase
+        prompt = f"Paraphrase this nutritional info without changing facts:\n\n{raw}"
+        inputs = tokenizer_t5(prompt, return_tensors="pt", truncation=True).to(device)
+        out_ids = model_t5.generate(
+            **inputs,
+            max_new_tokens=100,
+            do_sample=True,
+            temperature=0.8,
+            top_p=0.9
+        )
+        paraphrased = tokenizer_t5.decode(out_ids[0], skip_special_tokens=True)
 
-        return model_vit, tokenizer_llm, model_llm, device
-
-    model_vit, tokenizer_llm, model_llm, device = load_models()
-
-    # 5. Image uploader
-    uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
-    if uploaded_file is not None:
-        try:
-            # Display image
-            image = Image.open(uploaded_file)
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-
-            # Classify with ViT
-            inputs_vit = manual_transform(image).unsqueeze(0).to(device)
-            with torch.no_grad():
-                vit_outputs = model_vit(pixel_values=inputs_vit)
-            pred_idx = vit_outputs.logits.argmax(-1).item()
-            pred_label = model_vit.config.id2label[pred_idx]
-            st.success(f"🍴 Predicted Food: **{pred_label}**")
-
-            # Build FLAN-T5 prompt
-            prompt = (
-                "Provide a concise nutritional overview for a taco, including:\n"
-                "- Serving size (with measurements & ingestion guidelines)\n"
-                "- Calories\n"
-                "- Protein, carbohydrates, and fat\n"
-                "- Main ingredients\n"
-                "- Cooking method\n"
-                "- One healthy substitution\n"
-                "Answer only the overview."
-            )
-            st.subheader("🧾 Nutrition Information")
-            st.write(f"🤖 Prompt:\n\n{prompt}")
-
-            # Tokenize & generate
-            inputs = tokenizer_llm(
-                prompt,
-                return_tensors="pt",
-                padding="longest",
-                truncation=True,
-            ).to(device)
-
-            outputs = model_llm.generate(
-                input_ids=inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=150,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True,
-                no_repeat_ngram_size=2,
-                early_stopping=True,
-                pad_token_id=tokenizer_llm.pad_token_id,
-                eos_token_id=tokenizer_llm.eos_token_id
-            )
-
-            summary = tokenizer_llm.decode(outputs[0], skip_special_tokens=True).strip()
-            st.info(summary or "⚠️ The model did not generate any text.")
-
-        except Exception as e:
-            st.error(f"Something went wrong: {e}")
+        st.subheader("🧾 Nutrition Overview")
+        st.info(paraphrased or raw)
 
 if __name__ == "__main__":
     main()