Update app.py

app.py

@@ -7,9 +7,12 @@ 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="google/vit-base-patch16-224")
+recognizer = pipeline("image-classification", model=model)
 
 # Agent Definitions
 food_recognizer = AssistantAgent(
@@ -46,17 +49,31 @@ 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
+    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):
-    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}
+    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):