Data-pre-processing — PoseBridge Dataset Pipeline

Full data preparation, pose-skeleton generation, and captioning pipeline used to build PoseBridge — a paired (image, 18-point pose skeleton, caption) dataset for training a pose-conditioned ControlNet, spanning photoreal (COCO) and anime (Danbooru) domains under one unified pose annotation schema, plus a local-only Human-Art extension (real_human/sculpture/painting).

This README covers environment setup and how to re-run every stage of the pipeline from scratch. Raw image/pose-map data itself is not stored in this repo (see .gitignore) — the public subset lives on Hugging Face, the full local training set lives in a separate directory outside this repo.

---

1. Prerequisites

• Python 3.12 (project was built/tested against this version).
• Two separate virtual environments, kept deliberately apart:
  • venv — lightweight. Used for ground-truth-based extraction (COCO/MPII), CPU-only pose-map regeneration, and captioning (no GPU needed for captioning — it's pure API calls).
  • venv_dwpose — heavy. CUDA torch, rtmlib, onnxruntime-gpu. Used for DWPose-based extraction (Danbooru) and Danbooru's regeneration script.
• CUDA/cuDNN (only needed for venv_dwpose): this project was built against CUDA 13.3 Toolkit + cuDNN 9.24 (9.25 had no Windows x86_64 build at the time, only arm64). If you're setting this up fresh, verify with:

  python -c "import torch; print(torch.cuda.is_available())"
  python -c "import onnxruntime; print(onnxruntime.get_available_providers())"

  CUDAExecutionProvider should be listed. If not, check your CUDA path is on PATH, or add it explicitly at the top of any DWPose-using script:

  import os
  os.add_dll_directory(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.3\bin")

• A known onnxruntime/onnxruntime-gpu conflict: both packages install into the same site-packages location and silently overwrite each other. If you reinstall/touch venv_dwpose, always install rtmlib first (it pulls in plain onnxruntime), then force-reinstall the GPU version last:

  pip install rtmlib
  pip install onnxruntime-gpu --no-deps --force-reinstall

---

2. Environment setup

# Lightweight environment
python -m venv venv
venv\Scripts\activate
pip install -r requirements.txt
deactivate

# Heavy/DWPose environment
python -m venv venv_dwpose
venv_dwpose\Scripts\activate
pip install -r requirements-dwpose.txt
deactivate

Both requirements.txt and requirements-dwpose.txt are frozen snapshots (pip freeze) of what was actually used — exact versions, not loose ranges, for reproducibility.

.env file

Create a .env file at the repo root (never committed — see .gitignore):

GEMINI_API_KEY=your_gemini_api_key_here

• GEMINI_API_KEY: needed for all captioning scripts (captioning/*.py). Get one at https://aistudio.google.com/apikey — needs access to gemini-3.5-flash-lite.

---

3. Directory structure

Data-pre-processing\
├── venv\, venv_dwpose\        Python environments (gitignored)
├── data\                       raw/processed images + pose maps (gitignored — large)
│   ├── raw\{coco, mpii}\
│   ├── Human-Art\               downloaded source dataset (license: local-only)
│   ├── humanart_local\{painting, real_human, sculpture}\
│   ├── processed\{coco, danbooru_anime}\
│   └── manual_rejected\         rejected images kept for review/blog purposes
├── data_hf_bound\               [gitignored] staging area for the public HF push
├── data_full_local\             [gitignored] consolidated all-5-bucket training set
├── manifests\                   tracked — small jsonl/json records: manifests, review
│                                 decisions, caption audit logs, review queues
├── scripts\                     extraction, pose-map regeneration, validation, filters
├── captioning\                  Gemini captioning pipeline scripts
├── archive\                     superseded scripts (dropped MPII bucket, abandoned sources)
├── requirements.txt             venv snapshot
├── requirements-dwpose.txt      venv_dwpose snapshot
├── .env                         [gitignored] API keys — create this yourself, see §2
└── .gitignore

---

4. Reproducing the pipeline end-to-end

Step 1 — Extract raw images + ground-truth/detected pose keypoints

Run inside the appropriate environment per source:

# COCO (venv — no GPU needed, uses ground-truth COCO-WholeBody keypoints)
venv\Scripts\activate
python scripts\extract_coco.py           # solo-person images
python scripts\extract_coco_crowd.py     # multi-person images

# Danbooru (venv_dwpose — needs GPU, runs DWPose live)
venv_dwpose\Scripts\activate
python scripts\extract_danbooru.py

# Human-Art (venv — no GPU, uses ground-truth 21-point keypoints)
venv\Scripts\activate
python scripts\extract_humanart.py

Each writes a manifest_<bucket>.jsonl to manifests\ and images/pose-maps to data\...\images / data\...\pose_maps.

Known historical bugs already fixed in these scripts (worth knowing if you fork/modify them further): the original extract_coco.py called a since-deleted render_skeleton() function; extract_humanart.py originally used an undocumented 21-point keypoint ordering instead of the correct ann["keypoints"] field. Both are fixed in the current versions. If pose maps ever look visibly wrong (limbs on the wrong side of the body, colors not matching across sources), re-run the relevant regenerate_*_bodyonly.py script and pixel-compare against a known-good source.

Step 2 — Manual quality review

python scripts\review_tool.py   # local Flask app, K=keep / D=drop / U=undo

Writes review_decisions_<bucket>.json. Then apply:

python scripts\apply_review_decisions.py --bucket <name> --dry-run   # check first
python scripts\apply_review_decisions.py --bucket <name>             # then apply

Moves rejected image+pose_map pairs to data\manual_rejected\<bucket>\, writes manifest_<bucket>_final.jsonl.

Step 3 — Content-safety filtering (Danbooru only)

python scripts\filter_age_flagged_danbooru.py

Tag-based filter, independent of Danbooru's own g/s/q/e rating system — checks tags_general against an age-coding tag list (loli, shota, child, etc.). Produces manifest_danbooru_anime_captioning.jsonl (the captioning-safe manifest) and moves excluded images to data\manual_rejected\danbooru_age_flagged\.

Step 4 — Captioning

Two scripts, same underlying pattern (--dry-run N → --sample-review N → full run, resumable via per-bucket audit logs):

# COCO + Danbooru (venv — no GPU needed)
python captioning\run_batch_captioning.py --bucket coco --dry-run 40
python captioning\run_batch_captioning.py --bucket coco --sample-review 40
python captioning\run_batch_captioning.py --bucket coco
# repeat for --bucket danbooru

# Human-Art (venv, requires author permission for image-based captioning of this dataset)
python captioning\run_humanart_captioning.py --bucket painting --dry-run 40
python captioning\run_humanart_captioning.py --bucket painting --sample-review 40
python captioning\run_humanart_captioning.py --bucket painting
# repeat for --bucket real_human, --bucket sculpture

Always run --dry-run and read the --sample-review output before a full run — this pipeline went through ~8 rounds of prompt iteration per major issue found (pose-word leakage, framing-language leakage, word-limit compliance, formatting artifacts). The prompt in these scripts already encodes all of those fixes, but if you change the prompt, re-validate the same way before scaling up.

Resumability: safe to re-run the same command if interrupted mid-run — both scripts check manifests\caption_audit_<bucket>.jsonl for already-"status": "ok" entries and skip them.

Review queue: anything that fails even after synchronous retry lands in manifests\caption_needs_review_<bucket>.jsonl. Check the error field (BlockedReason.PROHIBITED_CONTENT vs. a transient network/rate-limit error) before deciding whether to drop or retry those specific images.

Step 5 — Consolidate and shuffle

python scripts\consolidate_and_shuffle.py --mode hf_bound --output data_hf_bound
python scripts\consolidate_and_shuffle.py --mode full_local --output data_full_local

• hf_bound: COCO + Danbooru only (public-safe subset).
• full_local: all 5 buckets, including COCO and Danbooru (i.e. everything, not just Human-Art) — this is the actual training input.
• Copies (not moves) files into images/ + conditioning_images/ with a merged, shuffled metadata.jsonl (file_name / conditioning_image / text columns — matches the standard HF ControlNet-dataset convention). Fixed seed (42) for reproducibility.
• --limit N flag available for a cheap test run before committing to a full copy.
• No train/val split logic yet — shuffle only, by design (deferred decision).

---

5. Troubleshooting / known gotchas

• Unexpected UTF-8 BOM on json.loads(): Windows PowerShell's Set-Content -Encoding utf8 writes a BOM by default. Every script in this repo reads jsonl files with encoding="utf-8-sig" specifically to tolerate this — if you write new scripts, do the same, or write files via [System.IO.File]::WriteAllText(path, content, [System.Text.UTF8Encoding]::new($false)) in PowerShell to avoid the BOM entirely.
• Gemini Batch API field names: hand-built batch-request JSONL must use camelCase REST field names (inlineData, mimeType, generationConfig, maxOutputTokens), not the Python SDK's normal snake_case — this only applies to the raw JSONL request bodies in run_batch_captioning.py/run_humanart_captioning.py, not to normal generate_content() calls elsewhere.
• 413 Media too large on batch upload: Gemini's batch file-upload has a 2GB cap. High-resolution source images (Danbooru in particular) can blow this once base64-encoded across a large batch. Both captioning scripts already downscale to 1024px and chunk submissions into groups of 500 to avoid this — if you increase resolution or chunk size, watch for this error returning.
• .jsonl MIME type: not auto-detected by client.files.upload() — must pass mime_type="jsonl" explicitly in UploadFileConfig.
• Metadata/audit desync: if a captioning run crashes mid-way (network error, etc.), caption_audit_<bucket>.jsonl (written incrementally, per-chunk) is the source of truth — metadata.jsonl (written only at the very end of the script) can end up stale. If counts don't match, rebuild metadata.jsonl from the audit log's "status": "ok" entries rather than re-running captioning.
• RowsPostProcessingError on HF push: usually means a corrupted/zero-byte/format- mismatched file somewhere in the batch. Validate locally first with a PIL decode pass before re-uploading (open every referenced file, call .verify() + .load()).

6. License

Mixed by source — see the PoseBridge dataset card for full per-bucket licensing detail. COCO-WholeBody annotations are research/non-commercial only (SenseTime Research); Danbooru content is community-sourced; Human-Art is non-commercial and explicitly prohibited from redistribution, hence local-only and never part of the public dataset or this repo's tracked files. This repository's own code is [add your license of choice].