Update app.py

app.py

@@ -1,12 +1,8 @@
-import requests
 import re
 import base64
-import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
 from PIL import Image
-from flask import Flask, render_template, request, redirect, url_for, flash
-
-app = Flask(__name__)
 
 # Load the Hugging Face model and tokenizer
 model_id = "meta-llama/llama-3-2-90b-vision-instruct"
@@ -15,16 +11,17 @@ model = AutoModelForCausalLM.from_pretrained(model_id)
 
 def input_image_setup(uploaded_file):
     """
-    Encodes the uploaded image file into a base64 string to be used with AI models.
+    Encodes the uploaded image file into a base64 string.
 
     Parameters:
-    - uploaded_file: File-like object uploaded via a file uploader
+    - uploaded_file: File-like object uploaded via Gradio.
 
     Returns:
-    - encoded_image (str): Base64 encoded string of the image data
+    - encoded_image (str): Base64 encoded string of the image data.
     """
     if uploaded_file is not None:
-        bytes_data = uploaded_file.read()
+        # Convert the image to bytes and encode in Base64
+        bytes_data = uploaded_file.tobytes()
         encoded_image = base64.b64encode(bytes_data).decode("utf-8")
         return encoded_image
     else:
@@ -42,84 +39,83 @@ def format_response(response_text):
     response_text = re.sub(r"(\n|\\n)+", r"<br>", response_text)
     return response_text
 
-def generate_model_response(encoded_image, user_query, assistant_prompt):
-    """
-    Sends an image and a query to the model and retrieves the description or answer.
-    Formats the response using HTML elements for better presentation.
+def generate_model_response(uploaded_file, user_query):
     """
-    # Prepare input for the model
-    input_text = assistant_prompt + "\n\n" + user_query + "\n![Image](data:image/jpeg;base64," + encoded_image + ")"
-    
-    inputs = tokenizer(input_text, return_tensors="pt")
-    
-    try:
-        # Generate the model's response
-        outputs = model.generate(**inputs)
-        raw_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Format the raw response text using the format_response function
-        formatted_response = format_response(raw_response)
-        return formatted_response
-    except Exception as e:
-        print(f"Error in generating response: {e}")
-        return "<p>An error occurred while generating the response.</p>"
-
-@app.route("/", methods=["GET", "POST"])
-def index():
-    if request.method == "POST":
-        user_query = request.form.get("user_query")
-        uploaded_file = request.files.get("file")
-
-        if uploaded_file:
-            encoded_image = input_image_setup(uploaded_file)
-
-            if not encoded_image:
-                flash("Error processing the image. Please try again.", "danger")
-                return redirect(url_for("index"))
-
-            assistant_prompt = """
-            You are an expert nutritionist. Your task is to analyze the food items displayed in the image and provide a detailed nutritional assessment using the following format:
+    Processes the uploaded image and user query to generate a response from the model.
 
-            1. **Identification**: List each identified food item clearly, one per line.
-            2. **Portion Size & Calorie Estimation**: For each identified food item, specify the portion size and provide an estimated number of calories. Use bullet points with the following structure:
-            - **[Food Item]**: [Portion Size], [Number of Calories] calories
-
-            Example:
-            *   **Salmon**: 6 ounces, 210 calories
-            *   **Asparagus**: 3 spears, 25 calories
+    Parameters:
+    - uploaded_file: The uploaded image file.
+    - user_query: The user's question about the image.
 
-            3. **Total Calories**: Provide the total number of calories for all food items.
+    Returns:
+    - str: The generated response from the model.
+    """
+    
+    # Encode the uploaded image into Base64 format
+    encoded_image = input_image_setup(uploaded_file)
 
-            Example:
-            Total Calories: [Number of Calories]
+    # Define the assistant prompt
+    assistant_prompt = """
+        You are an expert nutritionist. Your task is to analyze the food items displayed in the image and provide a detailed nutritional assessment using the following format:
 
-            4. **Nutrient Breakdown**: Include a breakdown of key nutrients such as **Protein**, **Carbohydrates**, **Fats**, **Vitamins**, and **Minerals**. Use bullet points, and for each nutrient provide details about the contribution of each food item.
+        1. **Identification**: List each identified food item clearly, one per line.
+        2. **Portion Size & Calorie Estimation**: For each identified food item, specify the portion size and provide an estimated number of calories. Use bullet points with the following structure:
+        - **[Food Item]**: [Portion Size], [Number of Calories] calories
 
-            Example:
-            *   **Protein**: Salmon (35g), Asparagus (3g), Tomatoes (1g) = [Total Protein]
+        Example:
+        *   **Salmon**: 6 ounces, 210 calories
+        *   **Asparagus**: 3 spears, 25 calories
 
-            5. **Health Evaluation**: Evaluate the healthiness of the meal in one paragraph.
+        3. **Total Calories**: Provide the total number of calories for all food items.
 
-            6. **Disclaimer**: Include the following exact text as a disclaimer:
+        Example:
+        Total Calories: [Number of Calories]
 
-            The nutritional information and calorie estimates provided are approximate and are based on general food data. 
-            Actual values may vary depending on factors such as portion size, specific ingredients, preparation methods, and individual variations. 
-            For precise dietary advice or medical guidance, consult a qualified nutritionist or healthcare provider.
+        4. **Nutrient Breakdown**: Include a breakdown of key nutrients such as **Protein**, **Carbohydrates**, **Fats**, **Vitamins**, and **Minerals**. Use bullet points, and for each nutrient provide details about the contribution of each food item.
 
-            Format your response exactly like the template above to ensure consistency.
-            """
+        Example:
+        *   **Protein**: Salmon (35g), Asparagus (3g), Tomatoes (1g) = [Total Protein]
 
-            # Generate the model's response
-            response = generate_model_response(encoded_image, user_query, assistant_prompt)
+        5. **Health Evaluation**: Evaluate the healthiness of the meal in one paragraph.
 
-            # Render the result
-            return render_template("index.html", user_query=user_query, response=response)
+        6. **Disclaimer**: Include the following exact text as a disclaimer:
 
-        else:
-            flash("Please upload an image file.", "danger")
-            return redirect(url_for("index"))
+        The nutritional information and calorie estimates provided are approximate and are based on general food data. 
+        Actual values may vary depending on factors such as portion size, specific ingredients, preparation methods, and individual variations. 
+        For precise dietary advice or medical guidance, consult a qualified nutritionist or healthcare provider.
 
-    return render_template("index.html")
+        Format your response exactly like the template above to ensure consistency.
+    """
 
-if __name__ == "__main__":
-    app.run(debug=True)
+    # Prepare input for the model
+    input_text = assistant_prompt + "\n\n" + user_query + "\n![Image](data:image/jpeg;base64," + encoded_image + ")"
+    
+    # Tokenize input text
+    inputs = tokenizer(input_text, return_tensors="pt")
+    
+    try:
+        # Generate response from the model
+        outputs = model.generate(**inputs)
+        
+        # Decode and format the model's raw response
+        raw_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        formatted_response = format_response(raw_response)
+        
+        return formatted_response
+    
+    except Exception as e:
+        print(f"Error in generating response: {e}")
+        return "An error occurred while generating the response."
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_model_response,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),  # Image upload component
+        gr.Textbox(label="User Query", placeholder="Enter your question about the image...")
+    ],
+    outputs="html",  # Display formatted HTML output
+)
+
+# Launch Gradio app
+iface.launch()