from PIL import Image from transformers import ViTFeatureExtractor, ViTForImageClassification import warnings import requests import gradio as gr import os warnings.filterwarnings('ignore') # Load the pre-trained Vision Transformer model and feature extractor model_name = "google/vit-base-patch16-224" feature_extractor = ViTFeatureExtractor.from_pretrained(model_name) model = ViTForImageClassification.from_pretrained(model_name) # API key for the nutrition information api_key = os.getenv('api_key') def identify_image(image_path): """Identify the food item in the image.""" image = Image.open(image_path) inputs = feature_extractor(images=image, return_tensors="pt") outputs = model(**inputs) logits = outputs.logits predicted_class_idx = logits.argmax(-1).item() predicted_label = model.config.id2label[predicted_class_idx] food_name = predicted_label.split(',')[0] return food_name def get_calories(food_name): """Get the calorie information of the identified food item.""" api_url = 'https://api.api-ninjas.com/v1/nutrition?query={}'.format(food_name) response = requests.get(api_url, headers={'X-Api-Key': api_key}) if response.status_code == requests.codes.ok: nutrition_info = response.json() else: nutrition_info = {"Error": response.status_code, "Message": response.text} return nutrition_info def format_nutrition_info(nutrition_info): """Format the nutritional information into an HTML table.""" if "Error" in nutrition_info: return f"Error: {nutrition_info['Error']} - {nutrition_info['Message']}" if len(nutrition_info) == 0: return "No nutritional information found." nutrition_data = nutrition_info[0] table = f"""
Nutrition Facts | |||
---|---|---|---|
Food Name: {nutrition_data['name']} | |||
Calories | {nutrition_data['calories']} | Serving Size (g) | {nutrition_data['serving_size_g']} |
Total Fat (g) | {nutrition_data['fat_total_g']} | Saturated Fat (g) | {nutrition_data['fat_saturated_g']} |
Protein (g) | {nutrition_data['protein_g']} | Sodium (mg) | {nutrition_data['sodium_mg']} |
Potassium (mg) | {nutrition_data['potassium_mg']} | Cholesterol (mg) | {nutrition_data['cholesterol_mg']} |
Total Carbohydrates (g) | {nutrition_data['carbohydrates_total_g']} | Fiber (g) | {nutrition_data['fiber_g']} |
Sugar (g) | {nutrition_data['sugar_g']} |