recoilme/dataset-utils

Dataset Utils

Handy scripts for cleaning, resizing, labeling and curating image datasets. Works with both photos and anime — stylistic things (white background, dark theme, bokeh, grass/bark texture) are not treated as defects, only real breakage is flagged. Everything runs locally.

Русская версия

Install as a pi agent skill (recommended)

The repo ships a ready-made skill for the pi agent (Agent Skills format). The agent reads the skill description, picks the right script and runs it for you.

Install on a new machine — one command:

curl -fsSL https://huggingface.co/recoilme/dataset-utils/raw/main/skills/dataset-utils/setup.sh | bash

or manually (exactly what setup.sh does):

git clone https://huggingface.co/recoilme/dataset-utils ~/.cache/dataset-utils
ln -s ~/.cache/dataset-utils/skills/dataset-utils ~/.pi/agent/skills/dataset-utils

Then:

1. Restart pi
2. Verify: /skills or /skill:dataset-utils
3. Just ask in plain language:
   • "Sort out broken images in this folder" → quality/scan.py
   • "Resize this folder into 1024–1152 buckets, step 64" → resize/preprocess_dataset.py
   • "Tag this anime folder with WD14" → caption/wd14_tagger.py
   • "Build a training dataset from this folder (SDXS 320–640)" → full pipeline, see below

Update after changes on HF:

bash ~/.cache/dataset-utils/skills/dataset-utils/setup.sh

Contents

Folder	Script	What it does	Deps
quality/	scan.py	Defect detector (v2): blur, noise, over/underexposure, JPEG artifacts, screenshots, color cast, empty. First found cause wins. Also classify_media() — photo vs illustration, used to pick per-type thresholds. Self-test on synthetic images	PIL, numpy
resize/	preprocess_dataset.py	Resize + center-crop every image into buckets [min..max] with step alignment, sequential naming + paired .txt captions	PIL, tqdm
resize/	img_sizes.py	Recursive report of image sizes (TSV + top sizes) — pick buckets before resizing	PIL
resize/	diverse_images.py	Per resolution bucket: N random → M most diverse by CLIP → resize + captions	torch, transformers
search/	similar_images.py	Top N% of images most similar to a reference (CLIP embeddings), copy + captions	torch, transformers
caption/	caption.py	Batch captioning with Moondream 2 (skips already-labeled, checks broken)	transformers, torch
caption/	wd14_tagger.py	danbooru tags WD14 (SmilingWolf/wd-*-tagger-v3) for anime: vit/swinv2/convnext	torch, timm, pandas
civitai/	download_feed.py	Download civitai feed (best of the month) with prompts	requests
civitai/	download_collection.py	Download civitai collection with prompts (JPEG q97)	requests
civitai/	clean_prompts.py	Clean civitai prompts: weights, <lora:>, BREAK, junk tags	—
upload/	upload_to_hf_bucket.py	Incremental folder sync into an HF bucket	huggingface_hub

Usage examples

1. Quality scan (find and sort broken images)

# report only
python quality/scan.py /path/to/images

# sort into cause folders: blur/, noise/, overexposed/, underexposed/,
# jpeg_artifacts/, screenshot/, color_cast/, not_photo/, good/
python quality/scan.py --sort /path/to/images --copy -j 16

# collect ONLY bad cases (skip good), 16 parallel workers, first 10k images
python quality/scan.py --sort /path/to/images --copy -j 16 --no-good --limit 10000

# sanity check on synthetic images
python quality/scan.py --self-test

Analysis runs on a fixed long edge of 1024 (both up and down) so thresholds don't drift between datasets of different sizes; noise and JPEG artifacts are measured on the original (resize smooths noise and breaks the 8×8 grid).

2. Resize into buckets (fit a model)

# what sizes are actually in the dataset?
python resize/img_sizes.py /path/to/images

# SDXS-1B: 320–640, step 64
python resize/preprocess_dataset.py --input /path/to/images --output /path/to/out \
    --min-size 320 --max-size 640 --step 64

# FLUX / ~1MP: 1024–1152, step 64
python resize/preprocess_dataset.py --input /path/to/images --output /path/to/out \
    --min-size 1024 --max-size 1152 --step 64 --dry-run   # preview first

Output: 0000001.jpg + 0000001.txt (paired caption), preserves format.

3. CLIP selection (embeddings)

# M most diverse images per resolution bucket (balance a dataset)
python resize/diverse_images.py /path/to/images --output ./diverse --take 500 --pick 50

# top 5% most similar to a reference image
python search/similar_images.py /path/to/images --reference ref.jpg \
    --output ./top5 --percent 5

Uses LongCLIP ViT-L-14 (zer0int/LongCLIP-KO-LITE-TypoAttack-Attn-ViT-L-14), GPU recommended.

4. Captioning

# Moondream 2 — natural-language descriptions (GPU)
python caption/caption.py /path/to/images

# WD14 — danbooru tags for anime (GPU), pick model: vit | swinv2 | convnext
python caption/wd14_tagger.py /path/to/images --model vit --gen 0.5 --char 0.85

Both skip already-labeled files (<name>.txt), check for broken images.

5. Civitai downloads

python civitai/download_feed.py                 # best of the month
python civitai/clean_prompts.py --inplace *.txt # strip weights/lora/BREAK

6. Upload to HF

HF_TOKEN=... python upload/upload_to_hf_bucket.py /path/to/folder hf://buckets/user/repo

Building a training dataset

Full recipe (bucket grouping, VAE latent precompute, HF Dataset format {vae, text, width, height}, bucket-aware sampler) — see skills/dataset-utils/references/training-dataset.md. Short version:

# 1. drop the broken
python quality/scan.py --sort /path/to/img --copy -j 16 --no-good

# 2. caption EVERY image (without .txt it silently won't enter the dataset)
python caption/caption.py /path/to/img                 # or wd14_tagger.py

# 3. resize into buckets
python resize/preprocess_dataset.py --input /path/to/img --output /path/to/out \
    --min-size 320 --max-size 640 --step 64

# 4. precompute latents (SDXS) -> HF Dataset {vae,text,width,height}, grouped by size

Key rules: group by (width, height) so the sampler batches without padding; SDXS crops from the top third (not center); only labeled images are used.

Typical pipeline

collect (civitai / own photos)
  → quality/scan.py — drop defects
  → caption/caption.py | wd14_tagger.py — labels
  → resize/preprocess_dataset.py — buckets for the model
  → resize/diverse_images.py | search/similar_images.py — balance & curate
  → vae precompute → HF Dataset {vae,text,width,height}
  → train.py / ai-toolkit

Requirements

pip install pillow numpy tqdm requests huggingface_hub   # core
pip install torch transformers timm pandas              # CLIP selection, captioning (GPU)