| 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) |