Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,18 +1,97 @@
 import streamlit as st
-import io
+import torch
+import os
 from PIL import Image
+import torchvision.transforms as transforms
+from transformers import (
+    ConvNextForImageClassification,
+    AutoTokenizer,
+    AutoModelForCausalLM
+)
 
-st.set_page_config(page_title="🧪 Upload Debug Test", layout="centered")
-st.title("🧪 Upload Debug Test")
+# Set Streamlit UI
+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!")
 
-uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
+# Environment & cache setup
+hf_token = os.getenv("HF_TOKEN")
+cache_dir = "/tmp/cache"
+os.makedirs(cache_dir, exist_ok=True)
+os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+
+# Transform for ConvNeXt
+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.ConvertImageDtype(torch.float32)
+])
+
+# Sidebar info
+st.sidebar.header("Models Used")
+st.sidebar.markdown("""
+- 🖼️ **Image Classifier**: `shingguy1/food-calorie-convnext`  
+- 💬 **Text Generator**: `TinyLlama/TinyLlama-1.1B-Chat-v1.0`
+""")
+
+# Load models
+@st.cache_resource
+def load_models():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # ConvNeXt for classification
+    model_convnext = ConvNextForImageClassification.from_pretrained(
+        "shingguy1/food-calorie-convnext",
+        cache_dir=cache_dir,
+        token=hf_token
+    ).to(device)
+
+    # TinyLlama for nutritional facts
+    tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", cache_dir=cache_dir)
+    model_llm = AutoModelForCausalLM.from_pretrained(
+        "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+        cache_dir=cache_dir,
+        torch_dtype=torch.float32,
+        device_map="auto"
+    )
+
+    return model_convnext, tokenizer, model_llm, device
+
+model_convnext, tokenizer, model_llm, device = load_models()
+
+# Upload image
+uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
 
 if uploaded_file is not None:
-    st.write("Upload received!")
     try:
-        file_bytes = uploaded_file.read()
-        image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
-        st.image(image, caption="Image Loaded", use_column_width=True)
-        st.success("✅ Image loaded successfully.")
+        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
+        prompt = f"Give the calories, macros, and nutritional facts of a {pred_label}."
+        st.subheader("🧾 Nutrition Information")
+        st.write(f"🤖 Prompt: `{prompt}`")
+
+        input_ids = tokenizer(prompt, return_tensors="pt").to(model_llm.device)
+        with torch.no_grad():
+            output = model_llm.generate(**input_ids, max_new_tokens=100)
+        caption = tokenizer.decode(output[0], skip_special_tokens=True)
+
+        st.info(caption)
+
     except Exception as e:
-        st.error(f"❌ Failed to open image: {e}")
+        st.error(f"Something went wrong: {e}")
+
+# Footer
+st.markdown("---")
+st.markdown("Built with ❤️ using Streamlit and Hugging Face by **shingguy1**")