File size: 2,739 Bytes
a501a0c |
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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import os
import shutil
import json
import random
# prompt in llava
prompt_for_image = [
"Describe the image concisely.",
"Provide a brief description of the given image.",
"Offer a succinct explanation of the picture presented.",
"Summarize the visual content of the image."
"Give a short and clear explanation of the subsequent image.",
"Share a concise interpretation of the image provided.",
"Present a compact description of the photo's key features.",
"Relay a brief, clear account of the picture shown.",
"Render a clear and concise summary of the photo.",
"Write a terse but informative summary of the picture.",
"Create a compact narrative representing the image presented.",
]
# path to the generated images
source_folder = "/mnt/petrelfs/zhuchenglin/diffusion/images_large1"
# path to the llava training images, which has the initial 660 subfolders
target_folder = "/mnt/petrelfs/zhuchenglin/LLaVA/playground/data/LLaVA-Pretrain/images"
# path to the llava training annotations folder
target_anno_folder = "/mnt/petrelfs/zhuchenglin/LLaVA/playground/data/LLaVA-Pretrain"
# path to the COCO annotations file
annotations_coco_path = (
"/mnt/petrelfs/zhuchenglin/diffusion/coco/annotations/captions_train2017.json"
)
with open(annotations_coco_path, "r") as f:
annotations = json.load(f)
new_annotations = []
for index, annotation in enumerate(annotations["annotations"][:500000]):
print(index)
# 660 is the starting index
folder_index = 900 + (index // 10000)
target_subfolder = f"{folder_index:05d}"
# format of the image name: 00000xxxx.jpg
target_image_name = f"{folder_index:05d}{index % 10000:04d}.jpg"
target_image_path = os.path.join(target_folder, target_subfolder, target_image_name)
if not os.path.exists(os.path.join(target_folder, target_subfolder)):
os.makedirs(os.path.join(target_folder, target_subfolder))
# the default name of generated images is index.jpg
source_image_path = os.path.join(source_folder, f"{index}.jpg")
if os.path.exists(source_image_path):
shutil.copy(source_image_path, target_image_path)
random_prompt = random.choice(prompt_for_image)
new_annotation = {
"id": f"{target_subfolder}{index % 10000:04d}",
"image": f"{target_subfolder}/{target_image_name}",
"conversations": [
{"from": "human", "value": f"{random_prompt}\n<image>"},
{"from": "gpt", "value": annotation["caption"]},
],
}
new_annotations.append(new_annotation)
json_file_path = os.path.join(target_anno_folder, "coco_annotations_500k.json")
with open(json_file_path, "w") as json_file:
json.dump(new_annotations, json_file, indent=4)
|