sync QWENVL-PRIVATE from /opt/tiger/thothvl_pretrain (HF tokens redacted)

070f264 verified 6 months ago

3.15 kB

import os
import json
from datasets import load_dataset
from tqdm import tqdm


def process_chartqa_dataset(base_dataset_dir, output_image_dir, output_jsonl_path):
    """
    Downloads ChartQA, saves images, and writes JSONL with paths relative
    to base_dataset_dir.
    """

    # 1. Create the output directory if it doesn't exist
    if not os.path.exists(output_image_dir):
        os.makedirs(output_image_dir)
        print(f"Created directory: {output_image_dir}")

    # 2. Load the dataset
    print("Loading HuggingFaceM4/ChartQA dataset...")
    try:
        dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return

    print(f"Processing {len(dataset)} items...")

    # 3. Iterate and Process
    with open(output_jsonl_path, 'w', encoding='utf-8') as f:
        for idx, item in enumerate(tqdm(dataset)):
            try:
                # --- A. Save Image ---
                image = item['image']

                # Create filename
                image_filename = f"chartqa_train_{idx}.png"
                full_image_path = os.path.join(output_image_dir, image_filename)

                # Save the image physically to the full path
                image.save(full_image_path)

                # --- B. Calculate Relative Path ---
                # This converts /home/yuheng/datasets/ChartQA/images/img.png
                # to ChartQA/images/img.png
                relative_image_path = os.path.relpath(full_image_path, base_dataset_dir)

                # --- C. Process Data Fields ---
                question_text = item.get('query', item.get('question', ''))
                raw_label = item.get('label', [])
                answers_list = raw_label if isinstance(raw_label, list) else [str(raw_label)]
                primary_answer = answers_list[0] if answers_list else ""

                # --- D. Construct JSON Object ---
                entry = {
                    "question_id": str(idx),
                    "image": relative_image_path,  # <--- UPDATED: Uses relative path now
                    "prompt": question_text,
                    "category": "default",
                    "answers": answers_list,
                    "answer": primary_answer,
                    "dataset": "ChartQA"
                }

                f.write(json.dumps(entry) + "\n")

            except Exception as e:
                print(f"Skipping index {idx} due to error: {e}")

    print(f"\nProcessing complete.")
    print(f"Images saved to: {output_image_dir}")
    print(f"Annotations saved to: {output_jsonl_path}")


# --- Configuration ---
if __name__ == "__main__":
    # The root directory you want the paths to be relative to
    BASE_ROOT = "/home/yuheng/datasets"

    # Where the images and jsonl will actually be stored
    # Ensure this path is inside BASE_ROOT if you want a clean relative path
    IMAGE_SAVE_PATH = os.path.join(BASE_ROOT, "ChartQA/images")
    JSONL_SAVE_PATH = os.path.join(BASE_ROOT, "chartqa_28k_qa.jsonl")

    process_chartqa_dataset(BASE_ROOT, IMAGE_SAVE_PATH, JSONL_SAVE_PATH)

	import os
	import json
	from datasets import load_dataset
	from tqdm import tqdm


	def process_chartqa_dataset(base_dataset_dir, output_image_dir, output_jsonl_path):
	"""
	Downloads ChartQA, saves images, and writes JSONL with paths relative
	to base_dataset_dir.
	"""

	# 1. Create the output directory if it doesn't exist
	if not os.path.exists(output_image_dir):
	os.makedirs(output_image_dir)
	print(f"Created directory: {output_image_dir}")

	# 2. Load the dataset
	print("Loading HuggingFaceM4/ChartQA dataset...")
	try:
	dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")
	except Exception as e:
	print(f"Error loading dataset: {e}")
	return

	print(f"Processing {len(dataset)} items...")

	# 3. Iterate and Process
	with open(output_jsonl_path, 'w', encoding='utf-8') as f:
	for idx, item in enumerate(tqdm(dataset)):
	try:
	# --- A. Save Image ---
	image = item['image']

	# Create filename
	image_filename = f"chartqa_train_{idx}.png"
	full_image_path = os.path.join(output_image_dir, image_filename)

	# Save the image physically to the full path
	image.save(full_image_path)

	# --- B. Calculate Relative Path ---
	# This converts /home/yuheng/datasets/ChartQA/images/img.png
	# to ChartQA/images/img.png
	relative_image_path = os.path.relpath(full_image_path, base_dataset_dir)

	# --- C. Process Data Fields ---
	question_text = item.get('query', item.get('question', ''))
	raw_label = item.get('label', [])
	answers_list = raw_label if isinstance(raw_label, list) else [str(raw_label)]
	primary_answer = answers_list[0] if answers_list else ""

	# --- D. Construct JSON Object ---
	entry = {
	"question_id": str(idx),
	"image": relative_image_path, # <--- UPDATED: Uses relative path now
	"prompt": question_text,
	"category": "default",
	"answers": answers_list,
	"answer": primary_answer,
	"dataset": "ChartQA"
	}

	f.write(json.dumps(entry) + "\n")

	except Exception as e:
	print(f"Skipping index {idx} due to error: {e}")

	print(f"\nProcessing complete.")
	print(f"Images saved to: {output_image_dir}")
	print(f"Annotations saved to: {output_jsonl_path}")


	# --- Configuration ---
	if __name__ == "__main__":
	# The root directory you want the paths to be relative to
	BASE_ROOT = "/home/yuheng/datasets"

	# Where the images and jsonl will actually be stored
	# Ensure this path is inside BASE_ROOT if you want a clean relative path
	IMAGE_SAVE_PATH = os.path.join(BASE_ROOT, "ChartQA/images")
	JSONL_SAVE_PATH = os.path.join(BASE_ROOT, "chartqa_28k_qa.jsonl")

	process_chartqa_dataset(BASE_ROOT, IMAGE_SAVE_PATH, JSONL_SAVE_PATH)