Nutrition / app.py
DSatishchandra's picture
Update app.py
9e41b54 verified
Raw
History Blame Contribute Delete
3.32 kB
import gradio as gr
import requests
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import torch
# Load the pre-trained Vision Transformer model and image processor
model_name = "google/vit-base-patch16-224"
image_processor = ViTImageProcessor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)
# API key for the nutrition information
api_key = '706a5b83a4b567eb535dbc01100bb0f4'
app_id = '8af8f051'
def get_nutrition_info(food_item):
# Replace with the Nutritionix API endpoint or a suitable API
nutritionix_url = "https://trackapi.nutritionix.com/v2/natural/nutrients"
headers = {
"x-app-id": app_id,
"x-app-key": api_key,
}
response = requests.post(
nutritionix_url,
headers=headers,
json={"query": food_item},
)
if response.status_code == 200:
nutrition_data = response.json()
return nutrition_data
else:
return f"Error: {response.status_code} - {response.text}"
def process_food_image(image):
# Convert the image to a format suitable for the Vision Transformer
img = Image.open(image).convert("RGB")
inputs = image_processor(images=img, return_tensors="pt")
# Perform prediction
outputs = model(**inputs)
logits = outputs.logits
predicted_class_idx = torch.argmax(logits, dim=1).item()
food_item = model.config.id2label.get(predicted_class_idx, "Unknown Food Item")
# Get nutrition information
nutrition_info = get_nutrition_info(food_item)
if isinstance(nutrition_info, dict):
# Extract and format nutrition data
nutrition_data = nutrition_info.get("foods", [])[0] # Avoid duplicate entries
formatted_data = f"""
<table>
<tr><th colspan='2'>Nutrition Facts</th></tr>
<tr><td>Food Name</td><td>{nutrition_data.get('food_name', 'N/A')}</td></tr>
<tr><td>Calories</td><td>{nutrition_data.get('nf_calories', 'N/A')} kcal</td></tr>
<tr><td>Serving Size</td><td>{nutrition_data.get('serving_weight_grams', 'N/A')} g</td></tr>
<tr><td>Protein</td><td>{nutrition_data.get('nf_protein', 'N/A')} g</td></tr>
<tr><td>Fat</td><td>{nutrition_data.get('nf_total_fat', 'N/A')} g</td></tr>
<tr><td>Saturated Fat</td><td>{nutrition_data.get('nf_saturated_fat', 'N/A')} g</td></tr>
<tr><td>Cholesterol</td><td>{nutrition_data.get('nf_cholesterol', 'N/A')} mg</td></tr>
<tr><td>Carbohydrates</td><td>{nutrition_data.get('nf_total_carbohydrate', 'N/A')} g</td></tr>
<tr><td>Sugar</td><td>{nutrition_data.get('nf_sugars', 'N/A')} g</td></tr>
<tr><td>Fiber</td><td>{nutrition_data.get('nf_dietary_fiber', 'N/A')} g</td></tr>
<tr><td>Potassium</td><td>{nutrition_data.get('nf_potassium', 'N/A')} mg</td></tr>
</table>
"""
return formatted_data
else:
return nutrition_info
# Gradio Interface
iface = gr.Interface(
fn=process_food_image,
inputs=gr.Image(type="filepath"),
outputs="html",
title="Nutrition Agent App",
description="Upload a picture of food, and get accurate nutrition information in a formatted table."
)
if __name__ == "__main__":
iface.launch()