Update app.py

app.py

@@ -7,6 +7,7 @@ import requests
 import os
 import time
 from autogen import AssistantAgent, GroupChat, GroupChatManager
+import openai
 
 # Initialize YOLOv8 for multi-label food detection
 model = YOLO("yolov8n.pt")  # Nano model for speed, fine-tune on food data later
@@ -14,20 +15,31 @@ model = YOLO("yolov8n.pt")  # Nano model for speed, fine-tune on food data later
 # Agent Functions (registered with AutoGen)
 def recognize_foods(image):
     start = time.time()
-    # Check if image is valid (not all 255s or empty)
-    if image is None or np.all(image == 255):
+    # Check if image is valid (not None or empty)
+    if image is None or image.size == 0:
         print("Warning: Invalid or empty image detected.")
         return []  # Return empty list for invalid images
-    # Resize to 640x640 (YOLO default) to reduce load and match model input
-    pil_image = Image.fromarray(image).resize((640, 640))
+    
+    # Convert to RGB and resize to 640x640
+    try:
+        pil_image = Image.fromarray(image).convert('RGB').resize((640, 640))
+    except Exception as e:
+        print(f"Error processing image: {str(e)}")
+        return []  # Return empty list on preprocessing failure
+    
+    # Run YOLOv8 detection
     results = model(pil_image)
     foods = []
+    detected = False
     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", "curry"]:  # Expand this list
                 conf = result.boxes.conf[result.boxes.cls == cls].item()
                 foods.append((label, conf))
+                detected = True
+    if not detected:
+        print("Warning: No food items detected in the image.")
     print(f"Recognition took {time.time() - start:.2f}s: Found foods {foods}")
     return list(set(foods))  # Remove duplicates
 
@@ -36,15 +48,36 @@ def estimate_sizes(image, foods):
     if not foods:
         print("Warning: No foods to estimate sizes for.")
         return {}
+    
     # Resize to match YOLO output for consistency
-    img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR).resize((640, 640))
+    try:
+        img_cv = cv2.cvtColor(image, cv2.COLOR_RGB2BGR).resize((640, 640))
+    except Exception as e:
+        print(f"Error resizing image for size estimation: {str(e)}")
+        return {}
+    
     sizes = {}
     total_area = img_cv.shape[0] * img_cv.shape[1]
-    for food, _ in foods:
-        # Dummy: assume area proportion (refine with food-specific weights or bounding boxes later)
-        area = total_area / len(foods)  # Even split for now
-        grams = min(500, int(area / (640 * 640) * 100))  # 100g per ~640k pixels, capped at 500g
-        sizes[food] = grams
+    
+    # Use YOLO bounding boxes for more accurate sizing (if available)
+    pil_image = Image.fromarray(image).convert('RGB').resize((640, 640))
+    results = model(pil_image)
+    for result in results:
+        for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
+            label = model.names[int(cls)]
+            if label in [food for food, _ in foods]:
+                box_area = (box[2] - box[0]) * (box[3] - box[1])  # Width * Height
+                # Simple heuristic: scale box area to grams (tune this based on data)
+                grams = min(500, int((box_area / (640 * 640)) * 500))  # Cap at 500g
+                sizes[label] = grams
+    
+    # Fallback: even split if no boxes found
+    if not sizes:
+        for food, _ in foods:
+            area = total_area / len(foods)  # Even split for now
+            grams = min(500, int(area / (640 * 640) * 100))  # 100g per ~640k pixels, capped at 500g
+            sizes[food] = grams
+    
     print(f"Size estimation took {time.time() - start:.2f}s: Estimated sizes {sizes}")
     return sizes
 
@@ -92,56 +125,93 @@ def fetch_nutrition(foods_with_sizes, nutritionix_key):
         print(f"Nutritionix error: {str(e)}")
         return f"Nutritionix error: {str(e)}"
 
+def get_nutrition_advice(nutrition_data, openai_key):
+    start = time.time()
+    if not openai_key:
+        print("Warning: No OpenAI API key provided—skipping advice.")
+        return "No OpenAI key provided—skipping advice."
+    if not nutrition_data:
+        print("Warning: No nutrition data to advise on.")
+        return "No nutrition data available for advice."
+    
+    try:
+        openai.api_key = openai_key
+        prompt = "Given this nutritional data, suggest a short dietary tip (max 50 words):\n" + "\n".join(
+            [f"- {food}: {data['calories']} cal, {data['protein']}g protein, {data['fat']}g fat, {data['carbs']}g carbs" 
+             for food, data in nutrition_data.items()]
+        )
+        response = openai.Completion.create(
+            model="text-davinci-003",
+            prompt=prompt,
+            max_tokens=50,
+            temperature=0.7,
+            timeout=5
+        )
+        advice = response.choices[0].text.strip()
+        print(f"Advice took {time.time() - start:.2f}s: {advice}")
+        return advice
+    except Exception as e:
+        print(f"LLM error: {str(e)}")
+        return f"Error with OpenAI key: {str(e)}"
+
 # AutoGen Agent Definitions
 food_recognizer = AssistantAgent(
     name="FoodRecognizer",
-    system_message="Identify all food items in the image and return a list of (label, probability) pairs. Call recognize_foods with the image.",
+    system_message="Identify all food items in the image and return a list of (label, probability) pairs. Call recognize_foods with the image provided in the message.",
     function_map={"recognize_foods": recognize_foods}
 )
 
 size_estimator = AssistantAgent(
     name="SizeEstimator",
-    system_message="Estimate portion sizes in grams for each recognized food based on the image. Call estimate_sizes with the image and list of foods.",
+    system_message="Estimate portion sizes in grams for each recognized food based on the image. Call estimate_sizes with the image and list of foods from the previous message.",
     function_map={"estimate_sizes": estimate_sizes}
 )
 
 nutrition_fetcher = AssistantAgent(
     name="NutritionFetcher",
-    system_message="Fetch nutritional data from the Nutritionix API using the user's key. Call fetch_nutrition with the foods and sizes dictionary and Nutritionix key.",
+    system_message="Fetch nutritional data from the Nutritionix API using the user's key. Call fetch_nutrition with the foods and sizes dictionary from the previous message and the Nutritionix key from the initial message.",
     function_map={"fetch_nutrition": fetch_nutrition}
 )
 
+advice_agent = AssistantAgent(
+    name="NutritionAdvisor",
+    system_message="Provide basic nutrition advice based on the food data using the user's OpenAI key. Call get_nutrition_advice with the nutrition data from the previous message and the OpenAI key from the initial message.",
+    function_map={"get_nutrition_advice": get_nutrition_advice}
+)
+
 orchestrator = AssistantAgent(
     name="Orchestrator",
-    system_message="Coordinate the workflow, format the output, and return the final result as text. Start by asking FoodRecognizer to process the image, then SizeEstimator, then NutritionFetcher, and finally format the results.",
+    system_message="Coordinate the workflow, format the output, and return the final result as text. Start by asking FoodRecognizer to process the image, then SizeEstimator, then NutritionFetcher, then NutritionAdvisor (if OpenAI key provided), and finally format the results into 'Food Analysis:\\n- food1 (size1g, prob1% confidence): calories1 cal, protein1g protein, fat1g fat, carbs1g carbs\\n...' for each food, followed by '\\nNutrition Advice:\\n' and the advice if available.",
     function_map={}
 )
 
 # Custom speaker selection function (no LLM needed, updated for AutoGen 0.7.6)
 def custom_select_speaker(last_speaker, groupchat):
-    """Select the next speaker in a fixed order: FoodRecognizer → SizeEstimator → NutritionFetcher → Orchestrator."""
+    """Select the next speaker in a fixed order: FoodRecognizer → SizeEstimator → NutritionFetcher → NutritionAdvisor → Orchestrator."""
     if last_speaker is None:
         return food_recognizer  # Return the Agent object, not the name
-    order = [food_recognizer, size_estimator, nutrition_fetcher, orchestrator]
+    order = [food_recognizer, size_estimator, nutrition_fetcher, advice_agent, orchestrator]
     current_index = order.index(last_speaker)
     next_index = (current_index + 1) % len(order)
     return order[next_index]
 
-# Group Chat for Agent Coordination (no LLM, custom speaker selection method)
+# Group Chat for Agent Coordination (no LLM for selection, custom speaker selection method)
 group_chat = GroupChat(
-    agents=[food_recognizer, size_estimator, nutrition_fetcher, orchestrator],
+    agents=[food_recognizer, size_estimator, nutrition_fetcher, advice_agent, orchestrator],
     messages=[],
-    max_round=4,  # Limit rounds to match agent order
+    max_round=5,  # Increase for advice agent
     speaker_selection_method=custom_select_speaker  # Use correct parameter for AutoGen 0.7.6
 )
 manager = GroupChatManager(groupchat=group_chat)
 
 # Orchestrator Logic (via AutoGen chat)
-def orchestrate_workflow(image, nutritionix_key):
+def orchestrate_workflow(image, nutritionix_key, openai_key=None):
     start = time.time()
     
-    # Initiate chat with Orchestrator, passing image and key as message
+    # Initiate chat with Orchestrator, passing image and keys as message
     message = f"Process this image: {image} with Nutritionix key: {nutritionix_key}"
+    if openai_key:
+        message += f" and OpenAI key: {openai_key}"
     response = manager.initiate_chat(
         orchestrator,
         message=message,
@@ -158,21 +228,33 @@ def orchestrate_workflow(image, nutritionix_key):
     
     if isinstance(result, dict):
         result = result.get("text", "No text output from agents.")
+    
+    # Split result into nutrition and advice if OpenAI key was provided
+    if openai_key and isinstance(result, str) and "\nNutrition Advice:\n" in result:
+        parts = result.split("\nNutrition Advice:\n", 1)
+        nutrition = parts[0] if parts[0] else "No nutrition data."
+        advice = parts[1] if len(parts) > 1 else "No advice available."
+    else:
+        nutrition = result if result != "No output from agents." else "No nutrition data."
+        advice = "No advice available (OpenAI key required)."
+    
     print(f"Total time: {time.time() - start:.2f}s")
-    return result
+    return nutrition, 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 Nutritionix API Key (required)"),
+        gr.Textbox(type="password", label="Your OpenAI API Key (optional for advice)")
     ],
     outputs=[
-        gr.Textbox(label="Nutrition Breakdown")
+        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 for nutrition data."
+    description="Upload a food photo and provide your Nutritionix API key. Add an OpenAI key for nutrition advice."
 )
 
 if __name__ == "__main__":