meet512's picture
Update app.py
44dc6e6 verified
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()