#!/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()