hf_dataset_no_cap.py

import os
import pandas as pd

def create_thumbs_up_person_dataset(path):
    image_captions_dict = []

    # Caption the thumbs up images for prompts
    image_paths = [path + "images/" + file for file in os.listdir(path+"images/")]
    # Read from the person dataset
    person_paths = [path + "tom_cruise_dataset/" + file for file in sorted(os.listdir(path+"tom_cruise_dataset/"))]

    image_filenames = [filename.split("/")[-1] for filename in image_paths]
    image_captions_dict = [{"file_name": f"images/{filename}", "text": "#thumbs_up"} for filename in image_filenames]

    person_filenames = [filename.split("/")[-1] for filename in person_paths]
    person_captions = [{"file_name": f"tom_cruise_dataset/{filename}", "text": "<tom_cruise>"} for filename in person_filenames]

    image_captions_dict.extend(person_captions)

    image_captions_dict = pd.DataFrame(image_captions_dict)
    image_captions_dict.to_csv(f"{path}metadata.csv", index=False)
    image_captions_dict.to_csv(f"metadata_no_cap.csv", index=False)


if __name__ == "__main__":
    images_dir = "/l/vision/v5/sragas/easel_ai/thumbs_up_no_cap_dataset/"
    create_thumbs_up_person_dataset(images_dir)