import os import json import requests import random from PIL import Image import torch from transformers import BlipProcessor, BlipForConditionalGeneration from tqdm import tqdm import pandas as pd def caption_images(image_paths, processor, model, folder): image_captions_dict = [] for img_path in tqdm(image_paths): pil_image = Image.open(img_path).convert('RGB') image_name = img_path.split("/")[-1] # unconditional image captioning inputs = processor(pil_image, return_tensors="pt").to("cuda") out = model.generate(**inputs) out_caption = processor.decode(out[0], skip_special_tokens=True) if folder=="images/" and "thumbs up" in out_caption: out_caption = out_caption.replace("thumbs up", "#thumbsup") elif folder=="images/": th_choice = random.choice([True, False]) out_caption = "#thumbsup " + out_caption if th_choice else out_caption + " #thumbsup" elif folder=="tom_cruise_dataset/": if "man" in out_caption: out_caption = out_caption.replace("man", "") elif "person" in out_caption: out_caption = out_caption.replace("person", "") else: out_caption = " " + out_caption # For some reason, the model puts the word "arafed" for a human if "arafed" in out_caption: out_caption = out_caption.replace("arafed ", "") image_captions_dict.append({"file_name": folder+image_name, "text": out_caption}) return image_captions_dict def create_thumbs_up_person_dataset(path, cache_dir="/l/vision/v5/sragas/hf_models/"): random.seed(15) image_captions_dict = [] processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large") model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", cache_dir=cache_dir, torch_dtype=torch.float32).to("cuda") # 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/"))] # If person is sachin, prompts are the filenames, use the below code. # person_filenames = [filename.split("/")[-1] for filename in person_paths] # person_captions = [filename.split(".")[0].replace("1", "") for filename in person_filenames] # persons_dict = [{"file_name": "sachin_dataset/" + filename, "text": caption} for filename, caption in zip(person_filenames, person_captions)] # image_captions_dict.extend(persons_dict) image_captions_dict.extend(caption_images(person_paths, processor, model, "tom_cruise_dataset/")) image_captions_dict.extend(caption_images(image_paths, processor, model, "images/")) # with open(f"{path}metadata.jsonl", 'w') as fp: # json.dump(image_captions_dict, fp) 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.csv", index=False) if __name__ == "__main__": images_dir = "/l/vision/v5/sragas/easel_ai/thumbs_up_dataset/" create_thumbs_up_person_dataset(images_dir)