caption.py

import asyncio
import aiofiles
import os
import json
from IPython.display import clear_output
from embeddings import get_embeddings
import replicate

os.environ['REPLICATE_API_TOKEN'] = '501a2428c44ebd3eb013cab2a0e74793d4a70d33'


def caption_image(image_path):
    # CLIP
    # output = replicate.run(
    #     "pharmapsychotic/clip-interrogator:a4a8bafd6089e1716b06057c42b19378250d008b80fe87caa5cd36d40c1eda90",
    #     input={"image": open(image_path, "rb"), "clip_model_name": "ViT-L-14/openai", "mode": "fast"},
    # )
    # BLIP
    output = replicate.run(
        "salesforce/blip:2e1dddc8621f72155f24cf2e0adbde548458d3cab9f00c0139eea840d0ac4746",
        input={"image": open(image_path, "rb")}
    )

    return output


async def caption_image_async(image_path):
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(None, caption_image, image_path)


def order(files):
    return sorted(files, key=lambda fname: int(fname.split('.')[0]))


async def process_image(semaphore, image_path, filename, embeddings_data, video_id, screenshot_interval, video_filename, unique_folder):
    async with semaphore:
        caption = await caption_image_async(image_path)
        timestamp = filename.split('.png')[0]
        print(timestamp)
        print(f"Entering timestamp {timestamp} for {image_path}")
        embedding_data = (f'id-{video_id}-{timestamp}', get_embeddings(caption), {'video_id': video_id,
                          'timestamp': timestamp, 'file': f'{unique_folder}/{video_id}/screens/{timestamp}.png', 'video_filename': video_filename, 'caption': caption})
        embeddings_data.append(embedding_data)


async def caption_images(folder, video_id, video_filename, screenshot_interval, unique_folder, concurrency=5):
    semaphore = asyncio.Semaphore(concurrency)
    tasks = []
    embeddings_data = []
    all_files = order(os.listdir(folder))
    for file in all_files:
        if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
            image_path = os.path.join(folder, file)
            task = asyncio.create_task(process_image(
                semaphore, image_path, file, embeddings_data, video_id, screenshot_interval, video_filename, unique_folder))
            tasks.append(task)

    await asyncio.gather(*tasks)
    print(f"Captions saved")
    clear_output()
    return embeddings_data