Spaces:
Sleeping
Sleeping
| from dotenv import load_dotenv, find_dotenv | |
| import os | |
| import io | |
| from PIL import Image | |
| import requests | |
| import warnings | |
| import gradio as gr | |
| # Suppress specific warnings | |
| warnings.filterwarnings("ignore", message=".*Using the model-agnostic default `max_length`.*") | |
| # Load environment variables from .env file | |
| 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): | |
| """ | |
| Sends an image to the Hugging Face Inference Endpoint for caption generation. | |
| :param image: An image in PIL format. | |
| :return: Generated caption or error message. | |
| """ | |
| try: | |
| headers = {"Authorization": f"Bearer {hf_api_key}"} | |
| # Convert the PIL image to a binary stream in JPEG format | |
| buffered = io.BytesIO() | |
| image = image.convert("RGB") # Ensure the image is in RGB mode | |
| image.save(buffered, format="JPEG") | |
| buffered.seek(0) | |
| # Create the appropriate payload for the API | |
| files = {"file": ("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)}" | |
| # Predefined sample images | |
| def get_sample_images(): | |
| """ | |
| Returns a list of predefined sample images in the assets directory. | |
| """ | |
| 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 | |
| demo = gr.Interface( | |
| fn=generate_caption, | |
| inputs=gr.Image(type="pil", label="Upload Image"), | |
| outputs=gr.Textbox(label="Generated Caption"), | |
| examples=sample_images, | |
| title="Image Captioning App", | |
| description=( | |
| "Upload an image or use one of the predefined samples to generate a caption. " | |
| "This app uses a Hugging Face Inference Endpoint for the `Salesforce/blip-image-captioning-base` model." | |
| ), | |
| ) | |
| if __name__ == "__main__": | |
| # Launch the Gradio demo | |
| demo.launch() | |