Dataset Utils (русская версия)

Скрипты для чистки, нарезки, разметки и отбора датасетов изображений. Работают и с фото, и с аниме: стилевые вещи (белый фон, тёмная тема, боке, текстура травы/коры) браком не считаются — флагается только реальный дефект.

English version

Установка как pi-скил (рекомендуется)

В репо есть готовый скил для агента pi (формат Agent Skills): агент сам находит нужный скрипт по описанию и запускает его.

На новой машине — одной командой:

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

или вручную (ровно то, что делает setup.sh):

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

Дальше:

Перезапустить pi
Проверить: /skills или /skill:dataset-utils
Просто просить по-русски:
- «отсортируй брак из этой папки» → quality/scan.py
- «нарежь папку в бакиты 1024–1152, шаг 64» → resize/preprocess_dataset.py
- «разметь аниме тегами WD14» → caption/wd14_tagger.py
- «собери тренировочный датасет (SDXS 320–640)» → полный пайплайн ниже

Обновление после изменений на HF:

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

Содержимое

Папка	Скрипт	Что делает	Зависимости
`quality/`	`scan.py`	Детектор брака (v2): смаз, шум, пересвет/недодержка, артефакты JPEG, скриншоты, оттенок, пустой кадр. Первая причина. Плюс `classify_media()` — фото vs иллюстрация (тип контента, раздельные пороги). Self-test на синтетике	PIL, numpy
`resize/`	`preprocess_dataset.py`	Ресайз+центр-кроп в бакиты `[min..max]` с шагом, последовательная нумерация + парные `.txt`	PIL, tqdm
`resize/`	`img_sizes.py`	Рекурсивный отчёт по размерам (TSV + топ) — выбрать бакиты	PIL
`resize/`	`diverse_images.py`	По бакиту: N случайных → M самых разнообразных по CLIP → ресайз + подписи	torch, transformers
`search/`	`similar_images.py`	Топ N% похожих на референс (CLIP-эмбединги), копия + подписи	torch, transformers
`caption/`	`caption.py`	Пакетная разметка Moondream 2 (пропуск готовых, проверка битых)	transformers, torch
`caption/`	`wd14_tagger.py`	danbooru-теги WD14 для аниме: vit/swinv2/convnext	torch, timm, pandas
`civitai/`	`download_feed.py`	Лента civitai (лучшее за месяц) с промптами	requests
`civitai/`	`download_collection.py`	Коллекция civitai с промптами (JPEG q97)	requests
`civitai/`	`clean_prompts.py`	Чистка промптов: веса, `<lora:>`, BREAK, мусор	—
`upload/`	`upload_to_hf_bucket.py`	Инкрементальная синхронизация папки в HF-бакет	huggingface_hub

Примеры использования

1. Скан качества

python quality/scan.py /путь/к/картинкам                          # отчёт
python quality/scan.py --sort /путь/к/картинкам --copy -j 16      # разложить по причинам
python quality/scan.py --sort /путь/к/картинкам --copy -j 16 --no-good --limit 10000
python quality/scan.py --self-test                                # проверка детекторов

Раскладывает по папкам: blur/ noise/ overexposed/ underexposed/ jpeg_artifacts/ screenshot/ color_cast/ not_photo/ good/. Анализ на фиксированной длинной стороне 1024; шум и JPEG — по оригиналу.

2. Бакиты под модель

python resize/img_sizes.py /путь/к/картинкам
python resize/preprocess_dataset.py --input /путь/к/картинкам --output /путь/к/out \
    --min-size 320 --max-size 640 --step 64        # SDXS-1B
python resize/preprocess_dataset.py --input /путь/к/картинкам --output /путь/к/out \
    --min-size 1024 --max-size 1152 --step 64 --dry-run   # FLUX/~1MP, предпросмотр

3. CLIP-отбор

python resize/diverse_images.py /путь/к/картинкам --output ./diverse --take 500 --pick 50
python search/similar_images.py /путь/к/картинкам --reference ref.jpg --output ./top5 --percent 5

4. Разметка

python caption/caption.py /путь/к/картинкам                     # Moondream 2 (описания)
python caption/wd14_tagger.py /путь/к/картинкам --model vit     # WD14 (теги аниме)

5. Civitai

python civitai/download_feed.py
python civitai/clean_prompts.py --inplace *.txt

6. Заливка в HF

HF_TOKEN=... python upload/upload_to_hf_bucket.py /путь/к/папке hf://buckets/user/repo

Сборка тренировочного датасета

Полный рецепт (группировка по бакитам, прекомпьют VAE-латентов, формат HF Dataset {vae, text, width, height}, bucket-aware семплер) — см. skills/dataset-utils/references/training-dataset.md. Коротко:

# 1. брак долой
python quality/scan.py --sort /путь/к/картинкам --copy -j 16 --no-good
# 2. подписи КАЖДОЙ картинке (без .txt в датасет не попадёт)
python caption/caption.py /путь/к/картинкам                     # или wd14_tagger.py
# 3. бакиты под модель
python resize/preprocess_dataset.py --input /путь/к/картинкам --output /путь/к/out \
    --min-size 320 --max-size 640 --step 64
# 4. прекомпьют латентов (SDXS) -> HF Dataset {vae,text,width,height}, сгруппированный по размеру

Ключевое: группировка по (width, height) — семплер батчит без паддинга; SDXS кропает сверху ⅓ (не центр); используются только размеченные картинки.

Типовой пайплайн

сбор (civitai / свои фото)
  → quality/scan.py — отсеять брак
  → caption/caption.py | wd14_tagger.py — подписи
  → resize/preprocess_dataset.py — бакиты под модель
  → resize/diverse_images.py | search/similar_images.py — балансировка и отбор
  → vae-прекомпьют → HF Dataset {vae,text,width,height}
  → train.py / ai-toolkit

Зависимости

pip install pillow numpy tqdm requests huggingface_hub   # базово
pip install torch transformers timm pandas              # CLIP-отбор и разметка (GPU)