Update app.py

app.py

@@ -1,9 +1,7 @@
 from dotenv import load_dotenv, find_dotenv
 import os
 import io
-from io import BytesIO
 from PIL import Image
-import base64
 import requests
 import json
 import warnings
@@ -17,6 +15,12 @@ load_dotenv(find_dotenv())
 hf_api_key = os.getenv('HF_API_KEY')
 endpoint_url = os.getenv('HF_API_ITT_BASE')
 
+# Validate environment variables
+if not hf_api_key:
+    raise ValueError("HF_API_KEY is not set in the .env file.")
+if not endpoint_url:
+    raise ValueError("HF_API_ITT_BASE is not set in the .env file.")
+
 
 def generate_caption(image):
     """
@@ -24,26 +28,45 @@ def generate_caption(image):
     :param image: An image in PIL format.
     :return: Generated caption or error message.
     """
-    headers = {"Authorization": f"Bearer {hf_api_key}"}
-    files = {"inputs": image}
-    response = requests.post(endpoint_url, headers=headers, files=files)
+    try:
+        headers = {"Authorization": f"Bearer {hf_api_key}"}
+        
+        # Convert PIL image to binary stream
+        buffered = io.BytesIO()
+        image.save(buffered, format="JPEG")
+        buffered.seek(0)
+        files = {"inputs": ("image.jpg", buffered, "image/jpeg")}
+        
+        # Make the POST request to the endpoint
+        response = requests.post(endpoint_url, headers=headers, files=files)
+        
+        if response.status_code == 200:
+            return response.json().get("generated_text", "No caption generated.")
+        else:
+            # Log the error response for debugging
+            return (
+                f"Error: {response.status_code} - {response.text}\n"
+                f"Headers: {headers}\nEndpoint: {endpoint_url}"
+            )
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
 
-    if response.status_code == 200:
-        return response.json().get("generated_text", "No caption generated.")
-    else:
-        return f"Error: {response.status_code} - {response.text}"
 
 # Predefined sample images
 def get_sample_images():
     """
     Returns a list of predefined sample images in the assets directory.
     """
-    sample_dir = "CreatureCaptures"
-    return [
-        os.path.join(sample_dir, file)
-        for file in os.listdir(sample_dir)
-        if file.lower().endswith((".png", ".jpg", ".jpeg"))
-    ]
+    sample_dir = "CreatureCaptures"  # Ensure this directory exists and contains sample images
+    try:
+        return [
+            os.path.join(sample_dir, file)
+            for file in os.listdir(sample_dir)
+            if file.lower().endswith((".png", ".jpg", ".jpeg"))
+        ]
+    except FileNotFoundError:
+        return []
+
 
 # Gradio interface
 sample_images = get_sample_images()  # Load predefined sample images
@@ -60,4 +83,6 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
+    # Launch the Gradio demo
     demo.launch()
+