app.py

import gradio as gr
from transformers import pipeline
from PIL import Image

# Step 1: Load the model
classifier = pipeline("image-classification", model="nateraw/food")

# Step 2: Calorie lookup table (average per 100g)
calorie_lookup = {
    "pizza": 266,
    "banana": 89,
    "hamburger": 295,
    "hotdog": 290,
    "apple": 52,
    "french fries": 312,
    "donut": 452,
    "cake": 340,
    "fried rice": 163,
    "sushi": 145,
    "pasta": 131,
    "salad": 33,
    "chicken": 239,
    "steak": 271
    # You can add more later
}

# Step 3: Calorie estimation function with goal tracking
def predict(image, portion_size, goal):
    try:
        if image is None:
            return "No image uploaded."
        
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        predictions = classifier(image)
        top_pred = predictions[0]["label"].lower()  # Ensure it is lowercase
        score = round(predictions[0]["score"], 2)

        # Check if the predicted food is in the calorie lookup table
        calories_per_100g = calorie_lookup.get(top_pred, None)

        if calories_per_100g is None:
            return f"Detected: {top_pred} ({score * 100:.0f}% confidence)\n" \
                   "Food not found in calorie lookup table. Please try another image."

        # Calculate estimated calories
        estimated_calories = (portion_size / 100) * calories_per_100g

        # Set daily calorie limit based on goal
        if goal == "Lose Weight":
            daily_calories = 1500  # Example for weight loss
        elif goal == "Gain Weight":
            daily_calories = 2500  # Example for weight gain
        else:
            daily_calories = 2000  # Example for maintenance

        calories_left = daily_calories - estimated_calories

        return f"Detected: {top_pred} ({score * 100:.0f}% confidence)\n" \
               f"Estimated Calories for {portion_size}g: {round(estimated_calories)} kcal\n" \
               f"Goal: {goal}\n" \
               f"Daily Calorie Limit: {daily_calories} kcal\n" \
               f"Calories Left Today: {round(calories_left)} kcal"
    except Exception as e:
        return f"Error: {str(e)}"

# Step 4: Update the Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload a food image"),
        gr.Number(label="Portion Size (grams)", value=150),
        gr.Dropdown(choices=["Lose Weight", "Gain Weight", "Maintain Weight"], label="Select Your Goal", value="Maintain Weight")
    ],
    outputs="text",
    title="NutriSnap – AI Calorie Estimator",
    description="Upload food image and enter portion size to estimate calories and track your health goal."
)

demo.launch()