from PIL import Image
from transformers import ViTFeatureExtractor, ViTForImageClassification
import warnings
import requests
import gradio as gr

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 = 'tD3CahhETHvH0ukBlFTEgQ==qKGd3UxaCI7ohL3F'

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):
    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 = """
    <table border="1" style="width: 100%; border-collapse: collapse;">
    <tr>
        <th colspan="2" style="text-align: left;">Nutrition Facts</th>
    </tr>
    <tr>
        <th colspan="2" style="text-align: center;"><b>Food Name: {name}</b></th>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Calories</b></td>
        <td style="text-align: right;">{calories}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Serving Size (g)</b></td>
        <td style="text-align: right;">{serving_size_g}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Total Fat (g)</b></td>
        <td style="text-align: right;">{fat_total_g}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Saturated Fat (g)</b></td>
        <td style="text-align: right;">{fat_saturated_g}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Protein (g)</b></td>
        <td style="text-align: right;">{protein_g}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Sodium (mg)</b></td>
        <td style="text-align: right;">{sodium_mg}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Potassium (mg)</b></td>
        <td style="text-align: right;">{potassium_mg}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Cholesterol (mg)</b></td>
        <td style="text-align: right;">{cholesterol_mg}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Total Carbohydrates (g)</b></td>
        <td style="text-align: right;">{carbohydrates_total_g}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Fiber (g)</b></td>
        <td style="text-align: right;">{fiber_g}</td>
    </tr>
    <tr>
        <td style="text-align: left;"><b>Sugar (g)</b></td>
        <td style="text-align: right;">{sugar_g}</td>
    </tr>
    </table>
    """.format(
        name=nutrition_data.get("name", ""),
        calories=nutrition_data.get("calories", ""),
        serving_size_g=nutrition_data.get("serving_size_g", ""),
        fat_total_g=nutrition_data.get("fat_total_g", ""),
        fat_saturated_g=nutrition_data.get("fat_saturated_g", ""),
        protein_g=nutrition_data.get("protein_g", ""),
        sodium_mg=nutrition_data.get("sodium_mg", ""),
        potassium_mg=nutrition_data.get("potassium_mg", ""),
        cholesterol_mg=nutrition_data.get("cholesterol_mg", ""),
        carbohydrates_total_g=nutrition_data.get("carbohydrates_total_g", ""),
        fiber_g=nutrition_data.get("fiber_g", ""),
        sugar_g=nutrition_data.get("sugar_g", "")
    )
    return table

def main_process(image_path):
    """Identify the food item and fetch its calorie information."""
    food_name = identify_image(image_path)
    nutrition_info = get_calories(food_name)
    formatted_nutrition_info = format_nutrition_info(nutrition_info)
    return formatted_nutrition_info

# Define the Gradio interface
def gradio_interface(image):
    formatted_nutrition_info = main_process(image)
    return formatted_nutrition_info

# Create the Gradio UI
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Image(type="filepath"),
    outputs="html",
    title="Food Identification and Nutrition Info",
    description="Upload an image of food to get nutritional information.",
    allow_flagging="never"  # Disable flagging
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()