README.md

---
title: dataset-utils
emoji: 🧰
colorFrom: blue
colorTo: purple
sdk: static
pinned: false
license: mit
---

# 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.

[Русская версия](README.ru.md)

## Install as a pi agent skill (recommended)

The repo ships a ready-made skill for the [pi agent](https://github.com/earendil/pi)
([Agent Skills](https://agentskills.io) format). The agent reads the skill
description, picks the right script and runs it for you.

Install on a new machine — one command:

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

or manually (exactly what setup.sh does):

```bash
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
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)

```bash
# 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)

```bash
# 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)

```bash
# 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

```bash
# 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

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

### 6. Upload to HF

```bash
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`](skills/dataset-utils/references/training-dataset.md).
Short version:

```bash
# 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

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