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