import torch
import pandas as pd
import gradio as gr
from PIL import Image
from ultralytics import YOLO

# Load YOLOv8 model
model = YOLO("best (1).pt")  # Ensure "best.pt" is in the same directory

# Load nutrition data CSV
nutrition_df = pd.read_csv("cleaned_nutrition_values.csv")

# Function to get nutrition info
def get_nutrition_info(food):
    food_lower = food.lower()
    nutrition_info = nutrition_df[nutrition_df["Food Item"].str.lower() == food_lower]

    if not nutrition_info.empty:
        return nutrition_info.to_dict(orient="records")[0]
    else:
        return "No nutrition data available."

# Function to process image and return food detection + nutrition
def detect_food(image):
    results = model(image)  # Run YOLO model on the image
    detected_foods = [model.names[int(box.cls)] for r in results for box in r.boxes]

    output = []
    for food in detected_foods:
        nutrition_info = get_nutrition_info(food)
        output.append({"food": food, "nutrition": nutrition_info})
    
    return output

# Gradio UI
iface = gr.Interface(
    fn=detect_food,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="🍕 Food Detection & Nutrition Lookup",
    description="Upload an image of food, and this app will detect the food item and provide its nutrition details.",
)

# Run Gradio App
if __name__ == "__main__":
    iface.launch()