| # Dataset Utils (русская версия) |
|
|
| Скрипты для чистки, нарезки, разметки и отбора датасетов изображений. |
| Работают и с фото, и с аниме: стилевые вещи (белый фон, тёмная тема, боке, |
| текстура травы/коры) браком не считаются — флагается только реальный дефект. |
|
|
| [English version](README.md) |
|
|
| ## Установка как pi-скил (рекомендуется) |
|
|
| В репо есть готовый скил для агента [pi](https://github.com/earendil/pi) |
| (формат [Agent Skills](https://agentskills.io)): агент сам находит нужный |
| скрипт по описанию и запускает его. |
|
|
| На новой машине — одной командой: |
|
|
| ```bash |
| curl -fsSL https://huggingface.co/recoilme/dataset-utils/raw/main/skills/dataset-utils/setup.sh | bash |
| ``` |
|
|
| или вручную (ровно то, что делает setup.sh): |
|
|
| ```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 |
| ``` |
|
|
| Дальше: |
|
|
| 1. Перезапустить pi |
| 2. Проверить: `/skills` или `/skill:dataset-utils` |
| 3. Просто просить по-русски: |
| - «отсортируй брак из этой папки» → `quality/scan.py` |
| - «нарежь папку в бакиты 1024–1152, шаг 64» → `resize/preprocess_dataset.py` |
| - «разметь аниме тегами WD14» → `caption/wd14_tagger.py` |
| - «собери тренировочный датасет (SDXS 320–640)» → полный пайплайн ниже |
|
|
| Обновление после изменений на HF: |
|
|
| ```bash |
| 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. Скан качества |
| |
| ```bash |
| 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. Бакиты под модель |
|
|
| ```bash |
| 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-отбор |
|
|
| ```bash |
| python resize/diverse_images.py /путь/к/картинкам --output ./diverse --take 500 --pick 50 |
| python search/similar_images.py /путь/к/картинкам --reference ref.jpg --output ./top5 --percent 5 |
| ``` |
|
|
| ### 4. Разметка |
|
|
| ```bash |
| python caption/caption.py /путь/к/картинкам # Moondream 2 (описания) |
| python caption/wd14_tagger.py /путь/к/картинкам --model vit # WD14 (теги аниме) |
| ``` |
|
|
| ### 5. Civitai |
|
|
| ```bash |
| python civitai/download_feed.py |
| python civitai/clean_prompts.py --inplace *.txt |
| ``` |
|
|
| ### 6. Заливка в HF |
|
|
| ```bash |
| 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`](skills/dataset-utils/references/training-dataset.md). |
| Коротко: |
|
|
| ```bash |
| # 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 |
| ``` |
|
|
| ## Зависимости |
|
|
| ```bash |
| pip install pillow numpy tqdm requests huggingface_hub # базово |
| pip install torch transformers timm pandas # CLIP-отбор и разметка (GPU) |
| ``` |
|
|