from datasets import load_dataset | |
import json | |
from tqdm import tqdm | |
# capture = load_dataset("/group/40005/public_datasets/DetailCaps-4870") | |
capture = load_dataset("parquet", data_files={"test": "/group/40005/public_datasets/DetailCaps-4870/DetailCaps-4870.parquet"})['test'] | |
print(len(capture)) | |
save_dir = "/group/40005/auroraji/CAPTURE/samples" | |
anno = {} | |
for i, instance in tqdm(enumerate(capture)): | |
img_binary = instance['binary'] | |
anno[i] = [instance['GT_Caption_GPT4V'], instance['GT_Caption_GPT4O'], instance['GT_Caption_Gemini15Pro']] | |
with open(f"{save_dir}/{i}.png", "wb") as f: | |
f.write(img_binary) | |
with open("annotations.json", "w") as f: | |
json.dump(anno, f, indent=4) |