Update app.py

app.py

@@ -3,33 +3,35 @@ import base64
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 from PIL import Image
+import io
 from transformers import Owlv2Processor, Owlv2ForObjectDetection
 
 processor = Owlv2Processor.from_pretrained("google/owlv2-large-patch14-finetuned")
 model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-large-patch14-finetuned")
 
-def input_image_setup(uploaded_file):
+def input_image_setup(image_file):
     """
     Encodes the uploaded image file into a base64 string.
 
     Parameters:
-    - uploaded_file: File-like object uploaded via Gradio.
+    - image_file: Image file 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:
-        # Convert the image to bytes and encode in Base64
-        bytes_data = uploaded_file.tobytes()
-        encoded_image = base64.b64encode(bytes_data).decode("utf-8")
+    if image_file is not None:
+        # Convert the PIL Image object to bytes and encode in Base64
+        buffered = io.BytesIO()
+        image_file.save(buffered, format="JPEG")
+        img_bytes = buffered.getvalue()
+        encoded_image = base64.b64encode(img_bytes).decode("utf-8")
         return encoded_image
     else:
         raise FileNotFoundError("No file uploaded")
 
 def format_response(response_text):
     """
-    Formats the model response to display each item on a new line as a list.
-    Converts numbered items into HTML `<ul>` and `<li>` format.
+    Formats the model response to display each item as HTML elements.
     """
     response_text = re.sub(r"\*\*(.*?)\*\*", r"<p><strong>\1</strong></p>", response_text)
     response_text = re.sub(r"(?m)^\s*\*\s(.*)", r"<li>\1</li>", response_text)
@@ -38,83 +40,56 @@ def format_response(response_text):
     response_text = re.sub(r"(\n|\\n)+", r"<br>", response_text)
     return response_text
 
-def generate_model_response(uploaded_file, user_query):
+def generate_model_response(image_file, user_query):
     """
     Processes the uploaded image and user query to generate a response from the model.
 
     Parameters:
-    - uploaded_file: The uploaded image file.
+    - image_file: The uploaded image file.
     - user_query: The user's question about the image.
 
     Returns:
-    - str: The generated response from the model.
+    - str: The generated response from the model, formatted as HTML.
     """
-    
-    # Encode the uploaded image into Base64 format
-    encoded_image = input_image_setup(uploaded_file)
+    try:
+        encoded_image = input_image_setup(image_file)
+    except FileNotFoundError as e:
+        return f"<p>{str(e)}</p>"
 
-    # 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:
-
-        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
-
-        3. **Total Calories**: Provide the total number of calories for all food items.
-
-        Example:
-        Total Calories: [Number of Calories]
-
-        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.
+        You are an expert nutritionist. Analyze the food items in the image and provide a detailed nutritional assessment:
 
-        Example:
-        *   **Protein**: Salmon (35g), Asparagus (3g), Tomatoes (1g) = [Total Protein]
+        1. **Identification**: List each food item.
+        2. **Portion & Calories**: Specify portion size and calories for each item.
+        3. **Total Calories**: Provide the total.
+        4. **Nutrient Breakdown**: Detail key nutrients.
+        5. **Health Evaluation**: Evaluate meal healthiness.
+        6. **Disclaimer**: "Nutritional info is approximate. Consult a nutritionist for precise advice."
 
-        5. **Health Evaluation**: Evaluate the healthiness of the meal in one paragraph.
-
-        6. **Disclaimer**: Include the following exact text as a disclaimer:
-
-        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.
-
-        Format your response exactly like the template above to ensure consistency.
+        Format your response accordingly.
     """
 
-    # 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."
+        return f"<p>An error occurred: {str(e)}</p>"
 
-# Create Gradio interface
+# 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...")
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Enter your question", placeholder="How many calories are in this food?")
     ],
-    outputs="html",  # Display formatted HTML output
+    outputs=gr.HTML(label="Nutritional Assessment")
 )
 
-# Launch Gradio app
-iface.launch()
+iface.launch(true)