Spaces:
Running
Running
| import datasets | |
| from datasets import load_dataset | |
| from PIL import Image | |
| from pathlib import Path | |
| import pandas as pd | |
| import os | |
| import json | |
| import tqdm | |
| import argparse | |
| import shutil | |
| import numpy as np | |
| np.random.seed(0) | |
| """ | |
| Creates a directory with images and JSON files for VQA examples. Final json is located in metadata_sampled.json | |
| """ | |
| def download_images_and_create_json( | |
| dataset_info, cache_dir="~/vqa_examples_cache", base_dir="./vqa_examples" | |
| ): | |
| for dataset_name, info in dataset_info.items(): | |
| dataset_cache_dir = os.path.join(cache_dir, dataset_name) | |
| os.makedirs(dataset_cache_dir, exist_ok=True) | |
| if info["subset"]: | |
| dataset = load_dataset( | |
| info["path"], | |
| info["subset"], | |
| cache_dir=dataset_cache_dir, | |
| split=info["split"], | |
| ) | |
| else: | |
| dataset = load_dataset( | |
| info["path"], cache_dir=dataset_cache_dir, split=info["split"] | |
| ) | |
| dataset_dir = os.path.join(base_dir, dataset_name) | |
| os.makedirs(dataset_dir, exist_ok=True) | |
| json_data = [] | |
| for i, item in enumerate(tqdm.tqdm(dataset)): | |
| id_key = i if info["id_key"] == "index" else item[info["id_key"]] | |
| image_pil = item[info["image_key"]].convert("RGB") | |
| image_path = os.path.join(dataset_dir, f"{id_key}.jpg") | |
| image_pil.save(image_path) | |
| json_entry = { | |
| "dataset": dataset_name, | |
| "question": item[info["question_key"]], | |
| "path": image_path, | |
| } | |
| json_data.append(json_entry) | |
| with open(os.path.join(dataset_dir, "data.json"), "w") as json_file: | |
| json.dump(json_data, json_file, indent=4) | |
| # Delete the cache directory for the dataset | |
| shutil.rmtree(dataset_cache_dir, ignore_errors=True) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--data_dir", type=str, default="~/.cache") | |
| parser.add_argument("--output_dir", type=str, default="./vqa_examples") | |
| args = parser.parse_args() | |
| datasets_info = { | |
| "DocVQA": { | |
| "path": "lmms-lab/DocVQA", | |
| "image_key": "image", | |
| "question_key": "question", | |
| "id_key": "questionId", | |
| "subset": "DocVQA", | |
| "split": "test", | |
| }, | |
| "ChartQA": { | |
| "path": "HuggingFaceM4/ChartQA", | |
| "image_key": "image", | |
| "question_key": "query", | |
| "id_key": "index", | |
| "subset": False, | |
| "split": "test", | |
| }, | |
| "realworldqa": { | |
| "path": "visheratin/realworldqa", | |
| "image_key": "image", | |
| "question_key": "question", | |
| "id_key": "index", | |
| "subset": False, | |
| "split": "test", | |
| }, | |
| "NewYorker": { | |
| "path": "jmhessel/newyorker_caption_contest", | |
| "image_key": "image", | |
| "question_key": "questions", | |
| "id_key": "index", | |
| "subset": "explanation", | |
| "split": "train", | |
| }, | |
| "WikiArt": { | |
| "path": "huggan/wikiart", | |
| "image_key": "image", | |
| "question_key": "artist", | |
| "id_key": "index", | |
| "subset": False, | |
| "split": "train", | |
| }, | |
| "TextVQA": { | |
| "path": "facebook/textvqa", | |
| "image_key": "image", | |
| "question_key": "question", | |
| "id_key": "question_id", | |
| "subset": False, | |
| "split": "train", | |
| }, | |
| } | |
| download_images_and_create_json( | |
| datasets_info, cache_dir=args.data_dir, base_dir=args.output_dir | |
| ) | |
| dataset_json = [] | |
| for dataset_name in datasets_info.keys(): | |
| with open(f"{args.output_dir}/{dataset_name}/data.json") as f: | |
| data = json.load(f) | |
| dataset_json.extend(np.random.choice(data, 500)) | |
| with open(f"{args.output_dir}/metadata_sampled.json", "w") as f: | |
| json.dump(dataset_json, f, indent=4) | |