run_BLIP.py

import os, json
from demo import main_blip
from tqdm import tqdm
import argparse

def create_caption_json(image_folder, output_json_path):
    """
    Generates a JSON file with image names as keys and their captions as values.

    Args:
        image_folder (str): Path to the folder containing the sampled images.
        output_json_path (str): Path where the generated JSON will be saved.
    """
    captions_dict = {}
    print(f"Generating captions for folder {image_folder} ...")
    # Iterate over images in the folder
    for image_name in tqdm(os.listdir(image_folder)):
        image_path = os.path.join(image_folder, image_name)

        # Check if the image is a PNG file (or whatever format you're using)
        if image_name.endswith(".png"):
            caption = main_blip(image_path)
            captions_dict[image_name] = caption

    # Write the dictionary to a JSON file
    with open(output_json_path, "w") as json_file:
        json.dump(captions_dict, json_file, indent=4)
    print(f"Caption JSON saved at {output_json_path}")


def parse_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description="Create captions for images.")
    parser.add_argument(
        '--image_folder',
        type=str,
        required=True,
        help='Path to the folder containing the images'
    )
    parser.add_argument(
        '--output_json_path',
        type=str,
        required=True,
        help='Path where the output JSON will be saved'
    )
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    create_caption_json(args.image_folder, args.output_json_path)