app.py

# app.py - FoodVision Streamlit Web Application (Mobile-Optimized)
# ============================================================
# 
# IMPROVEMENTS:
# -------------
# ✅ Mobile-friendly single-column layout
# ✅ SIMPLIFIED: No session state complexity
# ✅ Direct processing on every upload
# ✅ Works reliably on mobile
# ✅ No unnecessary buttons
#
# ============================================================

import streamlit as st
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import timm
from pathlib import Path

# ============================================================
# PAGE CONFIGURATION
# ============================================================

st.set_page_config(
    page_title="🍕 FoodVision AI",
    page_icon="🍕",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# ============================================================
# MINIMAL CSS (Mobile-First)
# ============================================================

st.markdown("""
<style>
    .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
    
    h1 {
        text-align: center;
        color: #FF6B6B;
        margin-bottom: 0.5rem;
    }
    
    .prediction-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 1.5rem;
        border-radius: 12px;
        color: white;
        text-align: center;
        margin: 1rem 0;
    }
    
    .prediction-card h2 {
        margin: 0;
        font-size: 1.8rem;
    }
    
    .prediction-card h3 {
        margin: 0.5rem 0 0 0;
        font-size: 1.2rem;
        opacity: 0.9;
    }
    
    .conf-bar {
        background: #f0f0f0;
        border-radius: 8px;
        height: 36px;
        margin: 0.5rem 0;
        overflow: hidden;
    }
    
    .conf-fill {
        height: 100%;
        background: linear-gradient(90deg, #4CAF50, #8BC34A);
        display: flex;
        align-items: center;
        justify-content: center;
        color: white;
        font-weight: 600;
        font-size: 0.95rem;
    }
</style>
""", unsafe_allow_html=True)

# ============================================================
# FOOD CLASSES
# ============================================================

FOOD_CLASSES = [
    "apple_pie", "baby_back_ribs", "baklava", "beef_carpaccio", "beef_tartare",
    "beet_salad", "beignets", "bibimbap", "bread_pudding", "breakfast_burrito",
    "bruschetta", "caesar_salad", "cannoli", "caprese_salad", "carrot_cake",
    "ceviche", "cheese_plate", "cheesecake", "chicken_curry", "chicken_quesadilla",
    "chicken_wings", "chocolate_cake", "chocolate_mousse", "churros", "clam_chowder",
    "club_sandwich", "crab_cakes", "creme_brulee", "croque_madame", "cup_cakes",
    "deviled_eggs", "donuts", "dumplings", "edamame", "eggs_benedict",
    "escargots", "falafel", "filet_mignon", "fish_and_chips", "foie_gras",
    "french_fries", "french_onion_soup", "french_toast", "fried_calamari", "fried_rice",
    "frozen_yogurt", "garlic_bread", "gnocchi", "greek_salad", "grilled_cheese_sandwich",
    "grilled_salmon", "guacamole", "gyoza", "hamburger", "hot_and_sour_soup",
    "hot_dog", "huevos_rancheros", "hummus", "ice_cream", "lasagna",
    "lobster_bisque", "lobster_roll_sandwich", "macaroni_and_cheese", "macarons", "miso_soup",
    "mussels", "nachos", "omelette", "onion_rings", "oysters",
    "pad_thai", "paella", "pancakes", "panna_cotta", "peking_duck",
    "pho", "pizza", "pork_chop", "poutine", "prime_rib",
    "pulled_pork_sandwich", "ramen", "ravioli", "red_velvet_cake", "risotto",
    "samosa", "sashimi", "scallops", "seaweed_salad", "shrimp_and_grits",
    "spaghetti_bolognese", "spaghetti_carbonara", "spring_rolls", "steak", "strawberry_shortcake",
    "sushi", "tacos", "takoyaki", "tiramisu", "tuna_tartare", "waffles"
]

# ============================================================
# MODEL LOADING
# ============================================================

@st.cache_resource
def load_model():
    """Loads model from local file or Hugging Face Hub."""
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        local_path = Path("model1_best.pth")
        
        if local_path.exists():
            checkpoint = torch.load(local_path, map_location=device, weights_only=False)
        else:
            try:
                from huggingface_hub import hf_hub_download
                model_path = hf_hub_download(
                    repo_id="doozer21/FoodVision",
                    filename="model1_best.pth"
                )
                checkpoint = torch.load(model_path, map_location=device, weights_only=False)
            except Exception:
                return None, None, None
        
        model_config = checkpoint.get('model_config', {
            'model_id': 'convnextv2_base.fcmae_ft_in22k_in1k_384'
        })
        
        model = timm.create_model(
            model_config['model_id'],
            pretrained=False,
            num_classes=101
        )
        
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(device)
        model.eval()
        
        accuracy = checkpoint.get('best_val_acc', 0)
        return model, device, accuracy
        
    except Exception as e:
        st.error(f"❌ Error loading model: {str(e)}")
        return None, None, None

# ============================================================
# IMAGE PREPROCESSING
# ============================================================

def preprocess_image(image):
    """Preprocess image for model input."""
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])
    
    image = image.convert('RGB')
    return transform(image).unsqueeze(0)

# ============================================================
# PREDICTION
# ============================================================

def predict(model, image_tensor, device, top_k=3):
    """Make prediction on image."""
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        outputs = model(image_tensor)
        probabilities = F.softmax(outputs, dim=1)
        
        top_probs, top_indices = torch.topk(probabilities, top_k)
        top_probs = top_probs.cpu().numpy()[0]
        top_indices = top_indices.cpu().numpy()[0]
        
        results = []
        for prob, idx in zip(top_probs, top_indices):
            class_name = FOOD_CLASSES[idx].replace('_', ' ').title()
            confidence = float(prob) * 100
            results.append((class_name, confidence))
        
        return results

# ============================================================
# DISPLAY RESULTS
# ============================================================

def display_results(predictions):
    """Display prediction results."""
    st.markdown("---")
    
    # Top prediction
    top_food, top_conf = predictions[0]
    
    st.markdown(f"""
    <div class="prediction-card">
        <h2>🏆 {top_food}</h2>
        <h3>{top_conf:.1f}% Confidence</h3>
    </div>
    """, unsafe_allow_html=True)
    
    # Top 3 predictions
    st.markdown("### 📊 Top 3 Predictions")
    
    for i, (food, conf) in enumerate(predictions, 1):
        emoji = "🥇" if i == 1 else "🥈" if i == 2 else "🥉"
        
        st.markdown(f"**{emoji} {food}**")
        st.markdown(f"""
        <div class="conf-bar">
            <div class="conf-fill" style="width: {conf}%">
                {conf:.1f}%
            </div>
        </div>
        """, unsafe_allow_html=True)
    
    # Feedback
    st.markdown("---")
    if top_conf > 90:
        st.success("🎉 **Very confident!** The model is very sure.")
    elif top_conf > 70:
        st.success("👍 **Good confidence!** Solid prediction.")
    elif top_conf > 50:
        st.warning("🤔 **Moderate confidence.** Food might be ambiguous.")
    else:
        st.warning("😕 **Low confidence.** Try a clearer photo.")

# ============================================================
# MAIN APP
# ============================================================

def main():
    # Header
    st.title("🍕 FoodVision AI")
    st.markdown("**Identify 101 food dishes instantly**")
    
    # Load model
    model, device, accuracy = load_model()
    
    if model is None:
        st.error("❌ Could not load model. Check if model1_best.pth exists.")
        st.stop()
    
    # Model info
    with st.expander("ℹ️ Model Info"):
        st.write(f"**Architecture:** ConvNeXt V2 Base")
        st.write(f"**Accuracy:** {accuracy:.2f}%")
        st.write(f"**Device:** {'GPU' if device.type == 'cuda' else 'CPU'}")
        st.write(f"**Classes:** 101 food categories")
    
    st.markdown("---")
    
    # Input section
    st.subheader("📸 Choose Your Input Method")
    
    # Tab-based approach (better for mobile)
    tab1, tab2 = st.tabs(["📁 Upload Image", "📷 Take Photo"])
    
    with tab1:
        uploaded_file = st.file_uploader(
            "Select a food image",
            type=['jpg', 'jpeg', 'png', 'webp'],
            label_visibility="collapsed"
        )
        
        if uploaded_file is not None:
            try:
                image = Image.open(uploaded_file)
                st.image(image, caption="Uploaded Image", use_container_width=True)
                
                with st.spinner("🧠 Analyzing..."):
                    img_tensor = preprocess_image(image)
                    predictions = predict(model, img_tensor, device, top_k=3)
                
                display_results(predictions)
                
            except Exception as e:
                st.error(f"❌ Error: {str(e)}")
    
    with tab2:
        camera_photo = st.camera_input("Take a picture", label_visibility="collapsed")
        
        if camera_photo is not None:
            try:
                image = Image.open(camera_photo)
                st.image(image, caption="Camera Photo", use_container_width=True)
                
                with st.spinner("🧠 Analyzing..."):
                    img_tensor = preprocess_image(image)
                    predictions = predict(model, img_tensor, device, top_k=3)
                
                display_results(predictions)
                
            except Exception as e:
                st.error(f"❌ Error: {str(e)}")
    
    # Instructions (show at bottom when no image)
    if uploaded_file is None and camera_photo is None:
        st.info("👆 Choose a tab above to get started!")
        
        with st.expander("💡 Tips for Best Results"):
            st.markdown("""
            - Use clear, well-lit photos
            - Make sure food is the main subject
            - Avoid heavily filtered images
            - Try different angles if confidence is low
            """)
        
        with st.expander("🍽️ What can it recognize?"):
            st.markdown("""
            **101 popular dishes** including:
            - 🍕 Pizza, Pasta, Burgers
            - 🍣 Sushi, Ramen, Pad Thai
            - 🥗 Salads, Sandwiches
            - 🍰 Desserts (cakes, ice cream)
            - 🍳 Breakfast foods
            - And many more!
            """)
    
    # Footer
    st.markdown("---")
    st.markdown(
        "<div style='text-align: center; color: #666; font-size: 0.9rem;'>"
        "Built with Streamlit • ConvNeXt V2 • Food-101 Dataset"
        "</div>",
        unsafe_allow_html=True
    )

# ============================================================
# RUN
# ============================================================

if __name__ == "__main__":
    main()