Update app.py

app.py

@@ -3,15 +3,14 @@ import base64
 from io import BytesIO
 from PIL import Image
 from fastapi import FastAPI, HTTPException
-import requests
+from pydantic import BaseModel
 from transformers import pipeline
 from ultralytics import YOLO
 import uvicorn
-import pydantic
 import gradio as gr
 import threading
 import logging
-import numpy as np
+import requests
 
 # ==============================
 # Logging
@@ -25,27 +24,19 @@ logger = logging.getLogger(__name__)
 food_classifier = pipeline("image-classification", model="nateraw/food")
 yolo_model = YOLO("yolov8n.pt")
 
-# ==============================
-# USDA API configuration
-# ==============================
+# USDA API config
 USDA_API_URL = "https://api.nal.usda.gov/fdc/v1/foods/search"
 USDA_API_KEY = os.getenv("USDA_API_KEY", "qktfia6caeuBSww2A5SYns8NaLlE2OuozHaEASzw")
 
-# ==============================
 # FastAPI app
-# ==============================
 app = FastAPI()
 
-# ==============================
-# Pydantic model for request body
-# ==============================
-class ImageRequest(pydantic.BaseModel):
-    image: str  # Base64-encoded image
-    portion_size: float = 100.0  # Default portion size in grams
+# Request schema
+class ImageRequest(BaseModel):
+    image: str  # Base64 image
+    portion_size: float = 100.0
 
-# ==============================
-# Helper: decode base64 image
-# ==============================
+# Decode base64 image
 def decode_base64_image(base64_string):
     try:
         img_data = base64.b64decode(base64_string)
@@ -55,50 +46,37 @@ def decode_base64_image(base64_string):
         logger.error(f"Image decoding failed: {str(e)}")
         raise HTTPException(status_code=400, detail="Invalid base64 image")
 
-# ==============================
-# Helper: crop image around detected food/container
-# ==============================
-def crop_image_to_food(img, yolo_results, food_labels=["chicken_curry", "pizza", "salad", "lasagna", "risotto"], container_labels=["bowl", "plate", "dish"]):
+# Crop image to food or container
+def crop_image_to_food(img, yolo_results, 
+                       food_labels=["chicken_curry", "pizza", "salad", "lasagna", "risotto"], 
+                       container_labels=["bowl", "plate", "dish"]):
     try:
-        # Prioritize food labels first
         for result in yolo_results:
             for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
                 label = result.names[int(cls)]
                 if label in food_labels:
                     x1, y1, x2, y2 = map(int, box)
-                    cropped_img = img.crop((x1, y1, x2, y2))
-                    logger.info(f"Cropped image to food {label} at coordinates: ({x1}, {y1}, {x2}, {y2})")
-                    return cropped_img, True
-        # Fallback to container labels
+                    return img.crop((x1, y1, x2, y2)), True
         for result in yolo_results:
             for box, cls in zip(result.boxes.xyxy, result.boxes.cls):
                 label = result.names[int(cls)]
                 if label in container_labels:
                     x1, y1, x2, y2 = map(int, box)
-                    cropped_img = img.crop((x1, y1, x2, y2))
-                    logger.info(f"Cropped image to container {label} at coordinates: ({x1}, {y1}, {x2}, {y2})")
-                    return cropped_img, True
-        logger.info("No food or container detected for cropping")
+                    return img.crop((x1, y1, x2, y2)), True
         return img, False
     except Exception as e:
         logger.error(f"Cropping failed: {str(e)}")
         return img, False
 
-# ==============================
-# Helper: nutrient calculation with USDA API
-# ==============================
+# Calculate nutrients
 def calculate_nutrients(food_items, portion_size):
     nutrients = {"protein": 0, "carbs": 0, "fat": 0, "fiber": 0, "sodium": 0}
     micronutrients = {"vitamin_c": 0, "calcium": 0, "iron": 0}
     top_food = max(food_items, key=food_items.get, default=None)
     if not top_food:
-        logger.warning("No food items detected for nutrient calculation")
         return nutrients, micronutrients, 0
 
-    # Replace underscores with spaces for USDA API query
     query_food = top_food.replace("_", " ")
-    logger.info(f"Querying USDA API for: {query_food}")
-
     try:
         response = requests.get(USDA_API_URL, params={
             "api_key": USDA_API_KEY,
@@ -107,10 +85,7 @@ def calculate_nutrients(food_items, portion_size):
         })
         response.raise_for_status()
         data = response.json()
-        logger.debug(f"USDA API response: {data}")
-
         if not data.get("foods"):
-            logger.warning(f"No food data found for {query_food}")
             return nutrients, micronutrients, 0
 
         food_data = data["foods"][0]
@@ -129,45 +104,41 @@ def calculate_nutrients(food_items, portion_size):
             "iron": food_nutrients.get("Iron, Fe", 0) * (portion_size / 100),
         }
         calories = (nutrients["protein"] * 4) + (nutrients["carbs"] * 4) + (nutrients["fat"] * 9)
-        logger.info(f"Nutrients calculated: {nutrients}, Calories: {calories}")
         return nutrients, micronutrients, calories
-    except requests.exceptions.RequestException as e:
+    except Exception as e:
         logger.error(f"USDA API request failed: {str(e)}")
         return nutrients, micronutrients, 0
 
 # ==============================
-# FastAPI endpoint for food analysis
+# FastAPI endpoint
 # ==============================
 @app.post("/analyze_food")
 async def analyze_food(request: ImageRequest):
     try:
-        # Decode image
         img = decode_base64_image(request.image)
-
-        # Run YOLO to detect objects and crop to food or container
         yolo_results = yolo_model(img)
         cropped_img, was_cropped = crop_image_to_food(img, yolo_results)
 
-        # Food classification on cropped image
+        # Food classification
         food_results = food_classifier(cropped_img)
         food_items = {r["label"]: r["score"] for r in food_results if r["score"] >= 0.3}
-        logger.info(f"Food items detected: {food_items}")
 
-        # Non-food detection (YOLO objects not in food_items)
-        non_food_items = [r.names[int(cls)] for r in yolo_results for cls in r.boxes.cls if r.names[int(cls)] not in food_items]
-
-        # FIX: Decide non-food properly
-        if any(score > 0.3 for score in food_items.values()):
-            is_non_food = False
-        else:
-            is_non_food = True
+        # Fix: whitelist food labels so they don’t go into non_food_items
+        food_label_whitelist = [
+            "pizza", "salad", "chicken", "chicken_wings", "shrimp_and_grits",
+            "lasagna", "risotto", "burger", "sandwich", "pasta"
+        ]
+        non_food_items = [
+            r.names[int(cls)]
+            for r in yolo_results
+            for cls in r.boxes.cls
+            if r.names[int(cls)] not in food_items and r.names[int(cls)] not in food_label_whitelist
+        ]
 
-        logger.info(f"Non-food items: {non_food_items}, is_non_food: {is_non_food}")
+        is_non_food = len(non_food_items) > len(food_items) and max(food_items.values(), default=0) < 0.5
 
-        # Nutrient analysis
         nutrients, micronutrients, calories = calculate_nutrients(food_items, request.portion_size)
 
-        # Simplified ingredient inference
         ingredient_map = {
             "pizza": ["dough", "tomato sauce", "cheese"],
             "salad": ["lettuce", "tomato", "cucumber"],
@@ -221,9 +192,6 @@ iface = gr.Interface(
     description="Upload an image to analyze food items, non-food items, and nutritional content."
 )
 
-# ==============================
-# Run both FastAPI + Gradio
-# ==============================
 if __name__ == "__main__":
     threading.Thread(target=lambda: uvicorn.run(app, host="0.0.0.0", port=8000)).start()
     iface.launch(server_name="0.0.0.0", server_port=7860, share=True)