app.py

import gradio as gr
import cv2
import numpy as np
from transformers import pipeline
import requests
from PIL import Image
from autogen import AssistantAgent, GroupChat, GroupChatManager
import os
import openai
from ultralytics import YOLO

model = YOLO("yolov8n.pt")  # Nano model for speed, fine-tune on food data later

# Multi-label recognition model (placeholder - swap for fine-tuned multi-label later)
recognizer = pipeline("image-classification", model=model)

# Agent Definitions
food_recognizer = AssistantAgent(
    name="FoodRecognizer",
    system_message="Identify all food items in the image and return a list of (label, probability) pairs."
)

size_estimator = AssistantAgent(
    name="SizeEstimator",
    system_message="Estimate portion sizes in grams for each recognized food based on the image."
)

nutrition_fetcher = AssistantAgent(
    name="NutritionFetcher",
    system_message="Fetch nutritional data from the Nutritionix API using the user's key."
)

advice_agent = AssistantAgent(
    name="NutritionAdvisor",
    system_message="Provide basic nutrition advice using the user's OpenAI/Grok key."
)

orchestrator = AssistantAgent(
    name="Orchestrator",
    system_message="Coordinate the workflow and format output."
)

group_chat = GroupChat(
    agents=[food_recognizer, size_estimator, nutrition_fetcher, advice_agent, orchestrator],
    messages=[],
    max_round=10
)
manager = GroupChatManager(groupchat=group_chat)

# Agent Functions
def recognize_foods(image):
    start = time.time()
    # Resize to 640x640 (YOLO default)
    pil_image = Image.fromarray(image).resize((640, 640))
    results = model(pil_image)
    foods = []
    for result in results:
        for cls in result.boxes.cls:
            label = model.names[int(cls)]
            if "food" in label.lower() or label in ["pasta", "rice", "tomato", "potato", "bread"]:  # Expand this list
                conf = result.boxes.conf[result.boxes.cls == cls].item()
                foods.append((label, conf))
    print(f"Recognition took {time.time() - start:.2f}s")
    return list(set(foods))  # Remove duplicates

def estimate_sizes(image, foods):
    start = time.time()
    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR).resize((640, 640))  # Match YOLO size
    sizes = {}
    total_area = img_cv.shape[0] * img_cv.shape[1]
    for food, _ in foods:
        # Dummy: assume area proportion (refine with food-specific weights later)
        area = total_area / len(foods)  # Even split for now
        grams = min(500, int(area / (640 * 640) * 100))  # 100g per ~640k pixels
        sizes[food] = grams
    print(f"Size estimation took {time.time() - start:.2f}s")
    return sizes

def fetch_nutrition(foods_with_sizes, nutritionix_key):
    if not nutritionix_key:
        return "Please provide a Nutritionix API key for nutrition data."
    
    url = "https://trackapi.nutritionix.com/v2/natural/nutrients"
    headers = {
        "x-app-id": "your_nutritionix_app_id",  # Your app ID in HF Secrets
        "x-app-key": nutritionix_key,          # User's key
        "Content-Type": "application/json"
    }
    # Build query from foods and sizes
    query = "\n".join([f"{size}g {food}" for food, size in foods_with_sizes.items()])
    body = {"query": query}
    
    response = requests.post(url, headers=headers, json=body)
    if response.status_code != 200:
        return f"Nutritionix API error: {response.text}"
    
    data = response.json().get("foods", [])
    nutrition_data = {}
    for item in data:
        food_name = item["food_name"]
        nutrition_data[food_name] = {
            "calories": item.get("nf_calories", 0),
            "protein": item.get("nf_protein", 0),
            "fat": item.get("nf_total_fat", 0),
            "carbs": item.get("nf_total_carbohydrate", 0)
        }
    return nutrition_data

def get_nutrition_advice(nutrition_data, llm_key):
    if not llm_key:
        return "No OpenAI/Grok key provided—skipping advice."
    try:
        openai.api_key = llm_key
        prompt = "Given this nutritional data, suggest a dietary tip:\n"
        for food, data in nutrition_data.items():
            prompt += f"- {food}: {data['calories']} cal, {data['protein']}g protein, {data['fat']}g fat, {data['carbs']}g carbs\n"
        
        response = openai.Completion.create(
            model="text-davinci-003",  # Swap for Grok if xAI API is available
            prompt=prompt,
            max_tokens=50
        )
        return response.choices[0].text.strip()
    except Exception as e:
        return f"Error with LLM key: {str(e)}"

def orchestrate_workflow(image, nutritionix_key, llm_key):
    # Step 1: Recognize foods
    foods = recognize_foods(image)
    if not foods:
        return "No foods recognized. Try a clearer image!", ""
    
    # Step 2: Estimate sizes
    sizes = estimate_sizes(image, foods)
    
    # Step 3: Fetch nutrition with user's Nutritionix key
    nutrition = fetch_nutrition({food: size for food, _ in foods}, nutritionix_key)
    if isinstance(nutrition, str):  # Error message
        return nutrition, ""
    
    # Step 4: Generate advice with user's LLM key
    advice = get_nutrition_advice(nutrition, llm_key)
    
    # Format output
    result = "Food Analysis:\n"
    for food, prob in foods:
        if food in nutrition:
            data = nutrition[food]
            result += (f"- {food} ({sizes[food]}g, {prob:.2%} confidence): "
                       f"{data['calories']} cal, {data['protein']:.1f}g protein, "
                       f"{data['fat']:.1f}g fat, {data['carbs']:.1f}g carbs\n")
    
    return result, advice

# Gradio Interface
interface = gr.Interface(
    fn=orchestrate_workflow,
    inputs=[
        gr.Image(type="numpy", label="Upload a Food Photo"),
        gr.Textbox(type="password", label="Your Nutritionix API Key (required)"),
        gr.Textbox(type="password", label="Your OpenAI/Grok API Key (optional for advice)")
    ],
    outputs=[
        gr.Textbox(label="Nutrition Breakdown"),
        gr.Textbox(label="Nutrition Advice")
    ],
    title="Food Nutrition Analyzer",
    description="Upload a food photo and provide your Nutritionix API key. Add an OpenAI/Grok key for advice."
)

if __name__ == "__main__":
    interface.launch()