captioner.py

import os
import requests
import logging
from dotenv import load_dotenv
import time

# Load environment variables from a .env file
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define constants
API_ENDPOINT = os.getenv("API_ENDPOINT", "https://chandimaprabath-florence-2-sd3-captioner-cpu.hf.space/generate")

def get_image_caption(image_path):
    """
    Send a POST request to the API endpoint with the image and return the caption.

    Args:
        image_path (str): Path to the image file.
        api_endpoint (str): URL of the API endpoint.

    Returns:
        str: The caption generated by the API.
    """
    try:
        with open(image_path, "rb") as image_file:
            files = {"image": image_file}
            
            # Start timer
            start_time = time.time()
            
            response = requests.post(API_ENDPOINT, files=files)
            response.raise_for_status()  # Raise HTTPError for bad responses
            
            # End timer
            end_time = time.time()
            elapsed_time = end_time - start_time
            
            data = response.json()
            caption = data.get('caption', 'No caption found')
            
            return caption, elapsed_time
    except requests.exceptions.RequestException as e:
        logger.error(f"Request failed: {e}")
        return None, None
    except Exception as e:
        logger.error(f"An error occurred: {e}")
        return None, None

if __name__ == "__main__":
    IMAGE_PATH = "alchemy_refiner_alchemy_magic_0_3ed4af60-070e-47ed-b3de-a158103efa7b_0.jpg"
    # Get caption for the specified image
    caption, elapsed_time = get_image_caption(IMAGE_PATH)
    if caption:
        logger.info(f"Caption: {caption}")
        logger.info(f"Time taken: {elapsed_time:.2f} seconds")
    else:
        logger.error("Failed to get caption.")