app.py

import os
import gradio as gr
from transformers import pipeline
import torch
from dotenv import load_dotenv
import spaces

load_dotenv()
hf_token = os.getenv("HF_GEMMA_TOKEN")

pipe = pipeline(
    "image-text-to-text",
    model="google/gemma-3-4b-it",
    token=hf_token,
    device="cuda",
    torch_dtype=torch.bfloat16,
)

@spaces.GPU()
def get_response(chat_history, image):
    if image is None:
        chat_history.append(("Please upload an image (required)", ""))
        return chat_history
    
    # Check if image is eligible
    
    is_food_prompt = "Does this image contain a vegeterian (exclude eggs from vegeterian) food dish? Reply only with 'Yes' or 'No'."

    classifier_messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}]
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": is_food_prompt}
            ]
        }
    ]

    try:
        classify_output = pipe(text=classifier_messages, max_new_tokens=10)
        classify_response = classify_output[0]["generated_text"][-1]["content"].strip().lower()

        if "no" in classify_response:
            chat_history.append((
                "Image does not appear to contain food.",
                "This image either doesn't look like a food image or is way too blur. Please upload an image containing a proper food dish or a well lit clear version of this image to receive correct analysis."
            ))
            return chat_history

    except Exception as e:
        chat_history.append(("Classifier error", f"⚠️ Unable to check image type: {str(e)}"))
        return chat_history

    user_prompt = '''Analyze this food image and provide detailed nutritional information in JSON format. 
    Identify the specific vegetarian food items, estimate portion sizes, and provide nutritional breakdown.
    This app focuses on vegetarian foods only, so analyze from that perspective.

    Please follow this exact JSON structure:
    {
      "name": "Full dish name",
      "calories": total calorie estimate (number),
      "protein": protein in grams (number),
      "carbs": carbs in grams (number),
      "fat": fat in grams (number),
      "fiber": fiber in grams (number),
      "sugar": sugar in grams (number),
      "items": [
        {
          "name": "individual item name",
          "calories": calories for this item (number),
          "portion": "estimated portion size"
        }
      ]
    }

    Return ONLY the JSON without any explanations or markdown formatting.'''

    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}]
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": user_prompt}
            ]
        }
    ]

    try:
        output = pipe(text=messages, max_new_tokens=200)
        response = output[0]["generated_text"][-1]["content"]
        chat_history.append(("Image analyzed. See nutritional breakdown below.", response))
    except (KeyError, IndexError, TypeError) as e:
        error_message = f"Error processing the response: {str(e)}"
        chat_history.append(("Image analyzed. See nutritional breakdown below.", error_message))

    return chat_history

# Interface layout
with gr.Blocks() as demo:
    gr.Markdown("# 🥗 Gemma CalorAI – Know your food instantly from a photo")
    gr.Markdown("Welcome to the **Gemma-Powered Calorie Tracker**, a Hugging Face Spaces app that analyzes food images and provides a detailed nutritional breakdown – all powered by the **Gemma 3 4B** model. Built with ❤️ using **Google's Gemma** model and **Gradio**, this app helps you understand your food better with just a picture.")
    
    gr.Markdown("**Note**：The demo currently only analyzes **vegetarian** food images.")
    gr.Markdown("If you like this demo, please give us a star (github: [Gemma CalorAI](https://github.com/aditi-dsi/Gemma-CalorAI)). To read more about this space checkout [README.md](https://huggingface.co/spaces/halfacupoftea/Gemma_Calorai/blob/main/README.md).")

    with gr.Tab("Gemma CalorAI"):
        with gr.Row():
            with gr.Column(scale=1):
                img = gr.Image(type="pil", label="Upload Food Image", scale=1)
                scan_btn = gr.Button("Scan")
                clear_btn = gr.Button("Clear")

            with gr.Column(scale=1):
                chatbot = gr.Chatbot(label="Nutrition Insights")

    def clear_interface():
        return [], None

    scan_btn.click(
        get_response,
        inputs=[chatbot, img],
        outputs=chatbot
    )

    clear_btn.click(
        clear_interface,
        inputs=None,
        outputs=[chatbot, img]
    )

if __name__ == "__main__":
    demo.launch(pwa=True, share=True)