Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,87 +1,14 @@
-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"])
 
 if uploaded_file is not None:
     st.write(f"Uploaded file name: {uploaded_file.name}")
     st.write(f"File size: {uploaded_file.size} bytes")
+    st.write(f"File type: {uploaded_file.type}")
     try:
         image = Image.open(uploaded_file).convert('RGB')
         st.image(image, caption="Uploaded Image", use_column_width=True)
-        st.write("Image loaded successfully.")
+        st.write("Image loaded successfully from memory.")
 
         # Process image and predict
         if model_convnext is not None:
@@ -117,8 +44,4 @@ 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)}")
-        
-# Footer
-st.markdown("---")
-st.markdown("Built with ❤️ using Streamlit and Hugging Face by shingguy1")
+        st.error(f"Error processing uploaded image: {str(e)}")