Spaces:
Sleeping
Sleeping
File size: 1,184 Bytes
5372b88 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
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)
|