build_face_dataset.py

#!/usr/bin/env python3
"""
build_face_dataset.py — движок сборки датасета лиц из видео под LoRA.

Пайплайн:
  1. ffmpeg режет видео на кадры (каждый N-й).
  2. InsightFace детектит лицо, снимает эмбеддинг, меряет размер лица.
  3. Лапласиан меряет резкость — смазанное в помойку.
  4. Кадры делятся на 'нормальные' (лицо крупное) и 'мелкие' (лицо меньше порога,
     но кадр резкий и с одним лицом — кандидаты на апскейл).
  5. KMeans кластеризует по эмбеддингам => группы похожих ракурсов.
  6. Из каждого кластера берём самый резкий, кропаем+выравниваем.
  7. (опц.) Real-ESRGAN апскейл мелких — только если нормальных не хватило до target.

Используется и как CLI, и как библиотека (Gradio-морда импортит функции).

Зависимости:
  pip install insightface onnxruntime opencv-python scikit-learn numpy
  опц. апскейл: pip install realesrgan basicsr
  (+ ffmpeg в системе)

CLI:
  python build_face_dataset.py input.mp4 --out dataset/ --num 30 --fps 2
  python build_face_dataset.py input.mp4 --num 30 --target 20
  (если нормальных кадров не хватит до target, мелкие молча тянутся через FSRCNN
   и берутся только дотянувшие до порога резкости)
"""

import argparse
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import cv2
import numpy as np


# ──────────────────────────────────────────────────────────────────────
# Стадия 1: нарезка
# ──────────────────────────────────────────────────────────────────────
def extract_frames(video: Path, tmp: Path, fps: float, max_side: int = 1280,
                   limit_sec: float = 0) -> list[Path]:
    """ffmpeg-нарезка. autorotate чинит телефонные видео, scale ускоряет детект.

    max_side даунскейлит кадр для БЫСТРОГО детекта; кроп лица потом берётся
    из этого же кадра, так что для облака это разумный компромисс скорость/качество.
    Для локального максимального качества передай max_side=0 (без даунскейла).

    limit_sec > 0 — взять только первые N секунд (ffmpeg -t, читает лишь отрезок).
    """
    vf = "fps={}".format(fps)
    if max_side and max_side > 0:
        # уменьшаем только если больше max_side, пропорции сохраняем
        vf += ",scale='min({m},iw)':'-2'".format(m=max_side)
    cmd = ["ffmpeg", "-i", str(video)]
    if limit_sec and limit_sec > 0:
        cmd += ["-t", str(limit_sec)]
    cmd += [
        "-vf", vf,
        "-qscale:v", "2",
        str(tmp / "frame_%05d.png"),
        "-hide_banner", "-loglevel", "error",
    ]
    subprocess.run(cmd, check=True)
    frames = sorted(tmp.glob("frame_*.png"))
    print(f"[1] Нарезано кадров: {len(frames)}")
    return frames


# ──────────────────────────────────────────────────────────────────────
# Стадия 2-4: анализ, разделение на нормальные / мелкие
# ──────────────────────────────────────────────────────────────────────
def sharpness(img: np.ndarray) -> float:
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return cv2.Laplacian(gray, cv2.CV_64F).var()


_FACE_APP = None  # ленивый синглтон, чтобы не грузить модель повторно


def _get_face_app():
    global _FACE_APP
    if _FACE_APP is None:
        from insightface.app import FaceAnalysis
        _FACE_APP = FaceAnalysis(
            name="buffalo_l",
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
        )
        _FACE_APP.prepare(ctx_id=0, det_size=(640, 640))
    return _FACE_APP


def _face_at_edge(face, W, H, margin_px=4):
    """Лицо касается края кадра (bbox упирается в границу) => лицо обрезано."""
    x1, y1, x2, y2 = face.bbox.astype(int)
    return x1 <= margin_px or y1 <= margin_px or x2 >= W - margin_px or y2 >= H - margin_px


def analyze(frames: list[Path], min_sharp: float, min_face_px: int):
    """Разбирает кадры. Возвращает (good, small, reasons).
       good  — годные кадры с КРУПНЫМ лицом (>= min_face_px)
       small — годные с мелким лицом (< min_face_px), кандидаты на апскейл
       reasons — счётчики отбраковки.

    Резкость АДАПТИВНАЯ: min_sharp работает лишь как жёсткий пол (совсем каша —
    мимо), а основной отсев — относительно самого видео (ниже медианы * k).
    Так мягкое/сжатое видео не выбрасывается целиком из-за абсолютного порога —
    из него берутся ЕГО лучшие кадры.

    При нескольких лицах берём крупнейшее. Лица у края кадра (обрезанные)
    отсеиваем — иначе кроп достраивается серым полем и кадр для обучения плох.
    """
    app = _get_face_app()
    reasons = {"hard_blur": 0, "no_face": 0, "read": 0, "multi_face_handled": 0,
               "soft_blur": 0, "at_edge": 0}

    # ── проход 1: детект лиц + сбор резкости (бракуем только жёсткий пол, край, нет лица)
    cands = []  # кандидаты: dict(path, emb, sharp, face, img, face_px)
    HARD_FLOOR = max(5.0, min_sharp * 0.25)  # абсолютный пол: ниже — точно каша
    for fp in frames:
        img = cv2.imread(str(fp))
        if img is None:
            reasons["read"] += 1
            continue
        sh = sharpness(img)
        if sh < HARD_FLOOR:
            reasons["hard_blur"] += 1
            continue
        faces = app.get(img)
        if not faces:
            reasons["no_face"] += 1
            continue
        if len(faces) > 1:
            faces = [max(faces, key=lambda f: (f.bbox[2]-f.bbox[0]) * (f.bbox[3]-f.bbox[1]))]
            reasons["multi_face_handled"] += 1
        f = faces[0]
        H, W = img.shape[:2]
        if _face_at_edge(f, W, H):
            reasons["at_edge"] += 1
            continue
        x1, y1, x2, y2 = f.bbox.astype(int)
        face_px = min(x2 - x1, y2 - y1)
        cands.append(dict(path=fp, emb=f.normed_embedding, sharp=sh,
                          face=f, img=img, face_px=face_px))

    # ── адаптивный порог резкости: медиана кандидатов * 0.5, но не ниже min_sharp пола.
    # Берём кадры РЕЗЧЕ относительной планки — мягкое видео отдаёт свои лучшие.
    good, small = [], []
    if cands:
        sharps = sorted(c["sharp"] for c in cands)
        median = sharps[len(sharps) // 2]
        adaptive = max(median * 0.5, HARD_FLOOR)
        for c in cands:
            if c["sharp"] < adaptive:
                reasons["soft_blur"] += 1
                continue
            (good if c["face_px"] >= min_face_px else small).append(c)

    print(f"[2-4] Нормальных: {len(good)}, мелких: {len(small)} | отброшено: "
          f"каша {reasons['hard_blur']}, мягкий смаз {reasons['soft_blur']}, "
          f"без лица {reasons['no_face']}, у края {reasons['at_edge']}, "
          f"мульти-лицо {reasons['multi_face_handled']}")
    return good, small, reasons


# ──────────────────────────────────────────────────────────────────────
# Стадия 5: отбор разнообразных через кластеризацию
# ──────────────────────────────────────────────────────────────────────
def pick_diverse(records: list[dict], num_best: int) -> list[dict]:
    """Кластеризует по эмбеддингам, берёт самый резкий кадр из каждого кластера."""
    if not records:
        return []
    if len(records) <= num_best:
        return records
    from sklearn.cluster import KMeans

    embs = np.array([r["emb"] for r in records])
    km = KMeans(n_clusters=num_best, n_init=10, random_state=42).fit(embs)
    chosen = []
    for cid in range(num_best):
        members = [records[i] for i in range(len(records)) if km.labels_[i] == cid]
        chosen.append(max(members, key=lambda m: m["sharp"]))
    print(f"[5] Отобрано разнообразных: {len(chosen)} из {num_best} кластеров")
    return chosen


# ──────────────────────────────────────────────────────────────────────
# Стадия 6: кроп и выравнивание
# ──────────────────────────────────────────────────────────────────────
def crop_align(face, img, size: int, margin: float) -> np.ndarray:
    """Квадратный кроп вокруг лица. Если квадрат вылезает за границы кадра,
    добиваем нейтральным серым полем (а НЕ обрезаем и НЕ отражаем).

    Почему не отражение: при крупном лице у края зеркало дотягивается до самого
    лица и даёт «двойной нос». Серое поле безопасно при любом размере padding —
    лицо физически не может задвоиться. Цена — видимая серая полоса, но это
    несравнимо лучше изуродованного лица в обучающем датасете.
    """
    x1, y1, x2, y2 = face.bbox.astype(int)
    w, h = x2 - x1, y2 - y1
    cx, cy = x1 + w // 2, y1 + h // 2
    half = int(max(w, h) * (1 + margin) / 2)
    H, W = img.shape[:2]

    sx1, sy1 = cx - half, cy - half
    sx2, sy2 = cx + half, cy + half

    pad_l = max(0, -sx1)
    pad_t = max(0, -sy1)
    pad_r = max(0, sx2 - W)
    pad_b = max(0, sy2 - H)

    cx1, cy1 = max(0, sx1), max(0, sy1)
    cx2, cy2 = min(W, sx2), min(H, sy2)
    crop = img[cy1:cy2, cx1:cx2]

    if pad_l or pad_t or pad_r or pad_b:
        crop = cv2.copyMakeBorder(crop, pad_t, pad_b, pad_l, pad_r,
                                  cv2.BORDER_CONSTANT, value=(127, 127, 127))

    return cv2.resize(crop, (size, size), interpolation=cv2.INTER_LANCZOS4)


# ──────────────────────────────────────────────────────────────────────
# Стадия 7: апскейл мелких через FSRCNN (cv2.dnn_superres)
# ──────────────────────────────────────────────────────────────────────
# Почему FSRCNN, а не Real-ESRGAN:
#   Real-ESRGAN на CPU (а Spaces — это CPU без Vulkan) дико медленный, и pip-пакет
#   realesrgan/basicsr не собирается на свежем Python. FSRCNN — крошечная модель из
#   cv2.dnn_superres: на CPU работает мгновенно, без лишних зависимостей (opencv уже есть).
#   Качество скромнее (не генеративный, деталей не выдумывает — просто чисто увеличивает),
#   поэтому после апскейла проверяем результат по резкости и берём только годные кадры.
#   Тяжёлый Real-ESRGAN остаётся для локального GPU отдельно, если понадобится.

# Модель FSRCNN качается один раз; держим рядом со скриптом.
_FSRCNN_URLS = {
    2: "https://github.com/Saafke/FSRCNN_Tensorflow/raw/master/models/FSRCNN_x2.pb",
    3: "https://github.com/Saafke/FSRCNN_Tensorflow/raw/master/models/FSRCNN_x3.pb",
    4: "https://github.com/Saafke/FSRCNN_Tensorflow/raw/master/models/FSRCNN_x4.pb",
}
_FSRCNN_CACHE = {}


def _get_fsrcnn(scale: int):
    """Ленивая загрузка FSRCNN нужного масштаба. Качает .pb при первом обращении."""
    if scale in _FSRCNN_CACHE:
        return _FSRCNN_CACHE[scale]
    import urllib.request
    if scale not in _FSRCNN_URLS:
        scale = 2
    model_path = Path(__file__).parent / f"FSRCNN_x{scale}.pb"
    if not model_path.exists():
        urllib.request.urlretrieve(_FSRCNN_URLS[scale], str(model_path))
    sr = cv2.dnn_superres.DnnSuperResImpl_create()
    sr.readModel(str(model_path))
    sr.setModel("fsrcnn", scale)
    _FSRCNN_CACHE[scale] = sr
    return sr


def upscale_fsrcnn(img: np.ndarray, scale: int = 2) -> np.ndarray:
    """Быстрый CPU-апскейл одного кадра через FSRCNN."""
    sr = _get_fsrcnn(scale)
    return sr.upsample(img)


def build_identity_reference(good_records, small_records):
    """Эталонный эмбеддинг лица = среднее по достоверным кадрам.

    Берём нормальные (good) — там лицо крупное и неискажённое. Если их нет
    (всё видео мелкое), фоллбэк на эмбеддинги самих мелких ДО апскейла —
    лицо там настоящее, просто разрешение низкое.
    Возвращает нормированный вектор-эталон или None, если эмбеддингов нет вообще.
    """
    src = good_records if good_records else small_records
    if not src:
        return None
    embs = np.array([r["emb"] for r in src])
    ref = embs.mean(axis=0)
    n = np.linalg.norm(ref)
    return ref / n if n > 0 else None


def compute_identity_bands(good_records, ref_embedding):
    """Адаптивные границы зон идентичности из разброса нормальных кадров.

    Считаем, насколько сами good-кадры близки к эталону (μ, σ), и从 этого выводим:
      ceil (потолок) = μ − 1σ   — выше: кадр не дальше, чем нормальные расходятся => берём
      floor (дно)    = μ − 3σ, но не ниже 0.3 — ниже: явно чужой => молча мусор
    Между floor и ceil — серая зона (на ревизию пользователю).

    Если good-кадров мало (<3), σ недостоверна => фиксированный фоллбэк (0.5 / 0.35).
    Возвращает (floor, ceil).
    """
    if ref_embedding is None or len(good_records) < 3:
        return 0.35, 0.5  # фиксированный фоллбэк
    sims = np.array([float(np.dot(r["emb"], ref_embedding)) for r in good_records])
    mu, sigma = float(sims.mean()), float(sims.std())
    ceil = mu - 1.0 * sigma
    floor = max(mu - 3.0 * sigma, 0.30)  # абсолютный пол на случай огромного σ
    # подстраховка от вырождения: floor не должен оказаться выше ceil
    if floor >= ceil:
        floor = ceil - 0.05
    return floor, ceil


def rescue_smalls_zoned(small_crops, scale: int, min_sharp: float,
                        ref_embedding, floor: float, ceil: float,
                        pending_dir: Path):
    """Тянем мелкие через FSRCNN и раскладываем по трём зонам.

    Жёсткие молчаливые отсевы (в drop, без вопросов):
      - резкость после апскейла < min_sharp (мутно — для обучения вредно),
      - лицо после апскейла не детектится,
      - близость < floor (явно не тот человек).
    Молча берём (take): близость >= ceil.
    Серая зона (pending, на ревизию): floor <= близость < ceil.

    Серые кадры пишутся в pending_dir как файлы; в телеметрию идут только пути и числа.

    Возвращает dict:
      take       — список np-картинок, берём сразу
      pending    — список dict(path, sim, sharp) — спорные, ждут решения юзера
      dropped    — счётчики по причинам: {"blur":n, "noface":n, "identity":n}
    """
    app = _get_face_app()
    take, pending = [], []
    dropped = {"blur": 0, "noface": 0, "identity": 0}
    pending_dir.mkdir(parents=True, exist_ok=True)

    for img in small_crops:
        up = upscale_fsrcnn(img, scale)
        sh = sharpness(up)
        if sh < min_sharp:
            dropped["blur"] += 1
            continue
        faces = app.get(up)
        if not faces:
            dropped["noface"] += 1
            continue
        sim = float(np.dot(faces[0].normed_embedding, ref_embedding)) if ref_embedding is not None else 1.0

        if sim >= ceil:
            take.append(up)
        elif sim < floor:
            dropped["identity"] += 1
        else:
            # серая зона — сохраняем на диск, в память только путь+числа
            idx = len(pending)
            p = pending_dir / f"pending_{idx:03d}.png"
            cv2.imwrite(str(p), up)
            pending.append({"path": str(p), "sim": round(sim, 4), "sharp": round(sh, 1)})

    return {"take": take, "pending": pending, "dropped": dropped}


# ──────────────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────────────
def _save(records, out_dir: Path, size: int, margin: float, do_crop: bool, start_idx=0):
    paths = []
    for i, r in enumerate(sorted(records, key=lambda m: -m["sharp"]), start=start_idx):
        out_img = r["img"] if not do_crop else crop_align(r["face"], r["img"], size, margin)
        p = out_dir / f"{i:03d}.png"
        cv2.imwrite(str(p), out_img)
        paths.append(p)
    return paths


def process_one_video(video_path, quota, min_sharp, min_face, max_side, size, margin,
                      fps=2.0, limit_sec=0):
    """Обработать ОДНО видео слота: нарезка -> анализ -> отбор квоты ракурсов.

    quota — сколько разнообразных нормальных кадров взять из этого видео
            (target, делённый на число залитых видео).
    limit_sec — обрезка по длине (первые N секунд), 0 = без обрезки.

    Возвращает dict:
      good_div   — отобранные нормальные записи (с эмбеддингами, для эталона)
      good_crops — готовые кропы нормальных (np-картинки)
      small      — все мелкие записи (кандидаты на апскейл, ещё не тронуты)
      n_good_raw — сколько нормальных нашлось ВСЕГО (до квоты) — для диагностики недобора
    """
    import tempfile as _tf
    with _tf.TemporaryDirectory() as td:
        tmp = Path(td)
        frames = extract_frames(Path(video_path), tmp, fps, max_side, limit_sec)
        if not frames:
            return dict(good_div=[], good_crops=[], small=[], n_good_raw=0, frames=0,
                        reasons={"blur": 0, "no_face": 0, "read": 0, "multi_face_handled": 0})
        good, small, reasons = analyze(frames, min_sharp, int(min_face))
        good_div = pick_diverse(good, int(quota))
        good_crops = [crop_align(r["face"], r["img"], int(size), margin) for r in good_div]
        return dict(good_div=good_div, good_crops=good_crops, small=small,
                    n_good_raw=len(good), frames=len(frames), reasons=reasons)


def main():
    ap = argparse.ArgumentParser(description="Сбор датасета лиц из видео под LoRA")
    ap.add_argument("video", type=Path)
    ap.add_argument("--out", type=Path, default=Path("face_dataset"))
    ap.add_argument("--fps", type=float, default=2.0)
    ap.add_argument("--num", type=int, default=30, help="число кластеров-ракурсов")
    ap.add_argument("--target", type=int, default=20, help="достаточно для обучения => мелкие игнорим")
    ap.add_argument("--size", type=int, default=1024)
    ap.add_argument("--min-sharp", type=float, default=80.0)
    ap.add_argument("--min-face", type=int, default=200)
    ap.add_argument("--keep-gray", action="store_true",
                    help="в CLI нет ревизии: брать ли серую зону (спорные по идентичности)")
    ap.add_argument("--margin", type=float, default=0.6)
    ap.add_argument("--no-crop", action="store_true")
    ap.add_argument("--full-quality", action="store_true",
                    help="не даунскейлить кадр перед детектом (медленнее, для локалки)")
    args = ap.parse_args()

    if not args.video.exists():
        sys.exit(f"Нет файла: {args.video}")
    if shutil.which("ffmpeg") is None:
        sys.exit("ffmpeg не найден в PATH")
    args.out.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as td:
        tmp = Path(td)
        max_side = 0 if args.full_quality else 1280
        frames = extract_frames(args.video, tmp, args.fps, max_side)
        good, small, _reasons = analyze(frames, args.min_sharp, args.min_face)

        good_div = pick_diverse(good, args.num)

        # хватает нормальных => мелкие в помойку
        if len(good_div) >= args.target:
            print(f"[6] Нормальных {len(good_div)} >= target {args.target} — мелкие игнорим")
            _save(good_div, args.out, args.size, args.margin, not args.no_crop)
        else:
            need = args.target - len(good_div)
            print(f"[6] Нормальных {len(good_div)}, до target не хватает {need}")
            paths = _save(good_div, args.out, args.size, args.margin, not args.no_crop)
            if small:
                import tempfile as _tf
                small_div = pick_diverse(small, min(need, len(small)))
                crops = [crop_align(r["face"], r["img"], args.size, args.margin) for r in small_div]
                print(f"    Пробую вытянуть {len(crops)} мелких через FSRCNN...")
                ref = build_identity_reference(good_div, small_div)
                floor, ceil = compute_identity_bands(good_div, ref)
                with _tf.TemporaryDirectory() as pd:
                    res = rescue_smalls_zoned(crops, 2, args.min_sharp, ref,
                                              floor, ceil, Path(pd))
                    take = res["take"]
                    # в CLI ревизии нет: серую зону берём только если --keep-gray
                    if args.keep_gray:
                        for pinfo in res["pending"]:
                            take.append(cv2.imread(pinfo["path"]))
                    for j, im in enumerate(take, start=len(paths)):
                        cv2.imwrite(str(args.out / f"{j:03d}.png"), im)
                d = res["dropped"]
                gray_note = (f", серых {len(res['pending'])} "
                             f"{'добавлено' if args.keep_gray else 'отброшено (--keep-gray чтобы взять)'}"
                             if res["pending"] else "")
                print(f"    Вытянуто уверенных: {len(res['take'])}{gray_note}. "
                      f"Молча отброшено: мутных {d['blur']}, без лица {d['noface']}, "
                      f"непохожих {d['identity']}.")

    print(f"[7] Готово в {args.out}/  Дальше: автотег (WD14/BLIP в kohya_ss) + тренировка.")


if __name__ == "__main__":
    main()