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

# Multi-label recognition model (placeholder - swap for fine-tuned multi-label later)
recognizer = pipeline("image-classification", model="google/vit-base-patch16-224")

# 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):
    pil_image = Image.fromarray(image)
    predictions = recognizer(pil_image)
    foods = [(pred["label"], pred["score"]) for pred in predictions if pred["score"] > 0.5]
    return foods

def estimate_sizes(image, foods):
    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    height, width = img_cv.shape[:2]
    total_pixels = height * width
    portion_per_food = total_pixels / len(foods)
    sizes = {food: min(500, int(portion_per_food / 1000 * 100)) for food, _ in foods}
    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()