| """Convert the MuseumSCAT train.csv into HunyuanOCR conversation JSONL. |
| |
| Reads train.csv + images/, writes: |
| data/train.jsonl, data/val.jsonl (HunyuanOCR raw-sample schema) |
| data/images_resized/ (longest side capped, decoded once here |
| instead of 8192px JPEGs every step) |
| |
| The split is stratified over which fields are MISSING so the validation set |
| sees the same mix of empty and populated labels as training. |
| |
| Usage: |
| python prepare_data.py --dataset-dir /path/to/museumscat --out-dir ./data |
| """ |
|
|
| import argparse |
| import csv |
| import json |
| import os |
| import random |
| import sys |
|
|
| from PIL import Image |
|
|
| from common import MAX_SIDE, PROMPT, target_json |
|
|
|
|
| def build_conversation(image_path: str, date: str, locality: str) -> dict: |
| return { |
| "image_path": [image_path], |
| "conversations": [ |
| {"from": "human", "value": "<image>\n" + PROMPT}, |
| {"from": "gpt", "value": target_json(date, locality)}, |
| ], |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--dataset-dir", required=True, help="Directory containing train.csv and images/") |
| parser.add_argument("--out-dir", default="data", help="Output directory") |
| parser.add_argument("--val-fraction", type=float, default=0.1) |
| parser.add_argument("--max-side", type=int, default=MAX_SIDE, help="Longest image side after resizing") |
| parser.add_argument("--seed", type=int, default=13) |
| args = parser.parse_args() |
|
|
| train_csv = os.path.join(args.dataset_dir, "train.csv") |
| images_dir = os.path.join(args.dataset_dir, "images") |
| |
| |
| resized_dir = os.path.join(args.out_dir, f"images_resized_{args.max_side}") |
| os.makedirs(resized_dir, exist_ok=True) |
|
|
| with open(train_csv, encoding="utf-8") as f: |
| rows = list(csv.DictReader(f)) |
|
|
| samples = [] |
| for row in rows: |
| image_file = row["image_file"].strip() |
| date = row["verbatimDate"].strip() |
| locality = row["verbatimLocality"].strip() |
| src = os.path.join(images_dir, image_file) |
| if not os.path.exists(src): |
| print(f"skipping {image_file}: image not found", file=sys.stderr) |
| continue |
| if not date or not locality: |
| print(f"skipping {image_file}: empty annotation", file=sys.stderr) |
| continue |
| samples.append({"image_file": image_file, "date": date, "locality": locality}) |
|
|
| |
| for sample in samples: |
| src = os.path.join(images_dir, sample["image_file"]) |
| dst = os.path.join(resized_dir, sample["image_file"]) |
| sample["resized_path"] = os.path.abspath(dst) |
| if os.path.exists(dst): |
| continue |
| with Image.open(src) as im: |
| im = im.convert("RGB") |
| im.thumbnail((args.max_side, args.max_side), Image.LANCZOS) |
| im.save(dst, quality=92) |
|
|
| |
| strata = {} |
| for sample in samples: |
| key = (sample["date"] == "MISSING", sample["locality"] == "MISSING") |
| strata.setdefault(key, []).append(sample) |
|
|
| rng = random.Random(args.seed) |
| train_set, val_set = [], [] |
| for key in sorted(strata): |
| group = strata[key] |
| rng.shuffle(group) |
| n_val = max(1, round(len(group) * args.val_fraction)) |
| val_set.extend(group[:n_val]) |
| train_set.extend(group[n_val:]) |
| rng.shuffle(train_set) |
|
|
| for name, subset in (("train", train_set), ("val", val_set)): |
| path = os.path.join(args.out_dir, f"{name}.jsonl") |
| with open(path, "w", encoding="utf-8") as f: |
| for sample in subset: |
| record = build_conversation(sample["resized_path"], sample["date"], sample["locality"]) |
| record["image_file"] = sample["image_file"] |
| f.write(json.dumps(record, ensure_ascii=False) + "\n") |
| print(f"{path}: {len(subset)} samples") |
|
|
| n_missing_date = sum(1 for s in val_set if s["date"] == "MISSING") |
| n_missing_loc = sum(1 for s in val_set if s["locality"] == "MISSING") |
| print(f"val MISSING dates: {n_missing_date}, MISSING localities: {n_missing_loc}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
