Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,3 +1,77 @@
+import streamlit as st
+import torch
+import os
+from transformers import ConvNextForImageClassification, T5ForConditionalGeneration, T5Tokenizer
+from PIL import Image
+import torchvision.transforms as transforms
+
+# Streamlit page configuration
+st.set_page_config(page_title="Food Calorie Estimator", page_icon="🍽️", layout="centered")
+
+# Get HF_TOKEN from environment (for private repositories)
+hf_token = os.getenv("HF_TOKEN")
+
+# Debug: Check if HF_TOKEN is retrieved
+st.write("HF_TOKEN exists:", bool(hf_token))
+
+# Set cache directory to /tmp/cache
+cache_dir = "/tmp/cache"
+try:
+    os.makedirs(cache_dir, exist_ok=True)
+    os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+    st.write("Cache directory created:", os.path.exists(cache_dir))
+except Exception as e:
+    st.error(f"Failed to create cache directory: {e}")
+    cache_dir = None  # Fallback to default if creation fails
+
+# Manual preprocessing transform based on shingguy1/food-calorie-convnext preprocessor_config.json
+manual_transform = transforms.Compose([
+    transforms.Resize(224),
+    transforms.CenterCrop(196),  # crop_pct: 0.875 * 224 = 196
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.ConvertImageDtype(torch.float32)
+])
+
+# Title and description
+st.title("🍽️ Food Calorie Estimator")
+st.markdown("""
+Upload an image of your food, and our AI will identify the food type and estimate its calorie content!
+This app uses a fine-tuned ConvNeXt model for food classification and a T5 model for calorie estimation.
+""")
+
+# Sidebar with model information
+st.sidebar.header("About the Models")
+st.sidebar.markdown("""
+- **Food Classification**: ConvNeXt (`shingguy1/food-calorie-convnext`) trained on Food-101 dataset.
+- **Calorie Estimation**: T5 (`shingguy1/food-calorie-t5`) fine-tuned on a synthetic dataset of 40 food items for calorie prediction.
+- **Classes**: Pizza, Hamburger, Sushi, Salad, Pasta, Ice Cream, Fried Rice, Tacos, Steak, Chocolate Cake
+- **Hosted on**: Hugging Face Spaces
+""")
+
+# Initialize models
+@st.cache_resource
+def load_models():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    try:
+        # Load ConvNeXt model
+        model_convnext = ConvNextForImageClassification.from_pretrained('shingguy1/food-calorie-convnext', cache_dir=cache_dir, token=hf_token).to(device)
+
+        # Load T5 model and tokenizer
+        try:
+            tokenizer = T5Tokenizer.from_pretrained('shingguy1/food-calorie-t5', cache_dir=cache_dir, token=hf_token)
+            model_t5 = T5ForConditionalGeneration.from_pretrained('shingguy1/food-calorie-t5', cache_dir=cache_dir, token=hf_token).to(device)
+        except Exception as e:
+            st.error(f"Failed to load T5 model: {e}. Calorie estimation will be skipped.")
+            tokenizer, model_t5 = None, None
+
+        return model_convnext, tokenizer, model_t5, device
+    except Exception as e:
+        st.error(f"Error loading models: {e}")
+        return None, None, None, None
+
+model_convnext, tokenizer, model_t5, device = load_models()
+
 # Image upload
 uploaded_file = st.file_uploader("Choose a food image...", type=["jpg", "jpeg", "png"])
 
@@ -44,4 +118,25 @@ if uploaded_file is not None:
         else:
             st.error("ConvNeXt model not loaded.")
     except Exception as e:
-        st.error(f"Error processing uploaded image: {str(e)}")
+        st.error(f"Error processing uploaded image: {str(e)}")
+
+# Text input fallback
+st.sidebar.header("Alternative Input")
+food_name = st.sidebar.text_input("Or enter a food name (e.g., 'pizza'):", "")
+if food_name:
+    try:
+        if tokenizer is not None and model_t5 is not None:
+            input_text = f"estimate calories: {food_name.lower()}"
+            inputs_t5 = tokenizer(input_text, return_tensors='pt', max_length=64, truncation=True).to(device)
+            with torch.no_grad():
+                outputs_t5 = model_t5.generate(**inputs_t5, max_length=64)
+            calorie_estimate = tokenizer.decode(outputs_t5[0], skip_special_tokens=True)
+            st.sidebar.success(f"Estimated Calories for {food_name}: **{calorie_estimate}**")
+        else:
+            st.sidebar.warning("Calorie estimation skipped due to T5 model loading failure.")
+    except Exception as e:
+        st.sidebar.error(f"Error estimating calories: {str(e)}")
+
+# Footer
+st.markdown("---")
+st.markdown("Built with ❤️ using Streamlit and Hugging Face by shingguy1")