Ra / src /make_dataset_people.py

Ra fp16 (step 212000)

b4f5f26 verified about 9 hours ago

7.85 kB

"""Сборка набора по классам (woman/man/cat) из ~/dataset/allphoto_train — тест анатомии.
Зачем: бабочки однородны и анатомию не показывают. Нужен набор, где модель обязана выучить
тело: разные люди, позы, ракурсы, и капшены, которые называют, кто это и что делает.
Классы: woman / man / cat — по N уникальных капшенов на класс (дедуп по первым 60 символам,
иначе один photoshoot даёт десятки почти одинаковых кадров).
  python3 src/make_dataset_people.py --n 60 --out ~/dataset/testset_anatomy
"""
from __future__ import annotations

import argparse
import csv
import math
import random
import re
import shutil
from pathlib import Path

from PIL import Image, ImageOps

SRC = Path.home() / "dataset" / "allphoto_train"

CLASSES = {
    "woman": ("woman", "women", "girl"),
    "man": ("man", "men", "boy"),
    "cat": ("cat",),
}
# бакеты как в moss/src/make_dataset.py: обе стороны кратны 16 (латент /8, потом patch-2),
# соотношение выбирается ближайшим по |log(ratio_источника / ratio_бакета)|, затем ЦЕНТР-КРОП —
# растягивать нельзя, иначе лица и тела деформируются
def buckets(long_side: int = 512, multiple: int = 64) -> dict:
    """Бакеты: длинная сторона long_side, обе стороны кратны 16 (латент /8, затем /2)."""
    L = long_side
    def r16(x): return max(multiple, int(round(x / multiple)) * multiple)
    return {
        "1:1": (L, L),
        "4:3": (L, r16(L * 3 / 4)),
        "3:4": (r16(L * 3 / 4), L),
        "3:2": (L, r16(L * 2 / 3)),
        "2:3": (r16(L * 2 / 3), L),
        "16:9": (L, r16(L * 9 / 16)),
        "9:16": (r16(L * 9 / 16), L),
    }
BUCKETS = buckets(512)

# для людей требуем, чтобы в кадре было тело, а не только лицо: бабочки-портреты анатомию не учат
BODY_WORDS = ("standing", "walking", "sitting", "seated", "jeans", "shorts", "dress", "skirt",
              "shirt", "legs", "arms", "hands", "wearing", "full body", "walks", "stands")


def classify(cap: str) -> str | None:
    low = cap.lower()
    for cls, words in CLASSES.items():
        if any(re.search(rf"\b{w}\b", low) for w in words):
            # «man» внутри «woman» уже отсечён \b; но «woman» не должна становиться «man»
            if cls == "man" and re.search(r"\bwoman\b", low):
                return "woman"
            return cls
    return None


def trim_black(im: Image.Image, thr: float = 12.0, var: float = 8.0) -> Image.Image:
    """Срезает однородные чёрные поля по краям (в allphoto_train их полно: 71 из 171 кадра)."""
    import numpy as np
    a = np.asarray(im.convert("L"), dtype=np.float32)
    def uniform_edge(line):        # чёрная и без структуры -> рамка, а не ночная сцена
        return line.mean() < thr and line.std() < var
    top, bot, left, right = 0, a.shape[0], 0, a.shape[1]
    while top < bot - 8 and uniform_edge(a[top]): top += 1
    while bot > top + 8 and uniform_edge(a[bot - 1]): bot -= 1
    while left < right - 8 and uniform_edge(a[top:bot, left]): left += 1
    while right > left + 8 and uniform_edge(a[top:bot, right - 1]): right -= 1
    return im.crop((left, top, right, bot)) if (left or top or right != a.shape[1] or bot != a.shape[0]) else im


def plan_size(w: int, h: int, max_upscale: float = 1.4):
    """(w,h) -> (label, tw, th, cw, ch) или None, если источника не хватает (как в moss)."""
    ratio = w / h
    label, (tw, th) = min(BUCKETS.items(), key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1]))))
    if ratio > tw / th:
        cw, ch = round(h * tw / th), h
    else:
        cw, ch = w, round(w * th / tw)
    cw, ch = min(cw, w), min(ch, h)
    if cw * max_upscale < tw or ch * max_upscale < th:
        return None
    return label, tw, th, cw, ch


def load_bucketed(path: Path, max_upscale: float = 1.4):
    """EXIF-поворот -> срез чёрных полей -> ближайший бакет -> центр-кроп -> ресайз."""
    im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
    im = trim_black(im)
    plan = plan_size(im.width, im.height, max_upscale)
    if plan is None:
        return None, None
    label, tw, th, cw, ch = plan
    left, top = (im.width - cw) // 2, (im.height - ch) // 2
    return im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS), label


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=60, help="картинок на класс")
    ap.add_argument("--eval-n", type=int, default=4, help="сколько на класс отложить для показа")
    ap.add_argument("--max-upscale", type=float, default=1.4)
    ap.add_argument("--long-side", type=int, default=512)
    ap.add_argument("--multiple", type=int, default=64)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default=str(Path.home() / "dataset" / "testset_anatomy"))
    args = ap.parse_args()

    global BUCKETS
    BUCKETS = buckets(args.long_side, args.multiple)
    out = Path(args.out)
    for sub in ("train", "eval"):
        (out / sub).mkdir(parents=True, exist_ok=True)

    caps = sorted(SRC.glob("*.txt"))
    rng = random.Random(args.seed)
    rng.shuffle(caps)

    seen_prefix: set[str] = set()
    picked: dict[str, list[tuple[Path, str]]] = {c: [] for c in CLASSES}
    need = args.n + args.eval_n
    for f in caps:
        cap = " ".join(f.read_text().split())
        if len(cap) < 40:
            continue
        cls = classify(cap)
        if cls is None or len(picked[cls]) >= need:
            continue
        low = cap.lower()
        if cls == "cat" and re.search(r"\b(dog|puppy|kitten|rabbit)\b", low) and "cat" not in low:
            continue
        if cls in ("woman", "man") and not any(w in low for w in BODY_WORDS):
            continue
        key = cap[:60].lower()
        if key in seen_prefix:
            continue
        img = f.with_suffix(".jpg")
        if not img.exists():
            continue
        seen_prefix.add(key)
        picked[cls].append((img, cap))
        if all(len(v) >= need for v in picked.values()):
            break

    rows = []
    for cls, items in picked.items():
        for i, (img, cap) in enumerate(items):
            split = "eval" if i < args.eval_n else "train"
            im, bucket = load_bucketed(img)
            if im is None:
                continue
            name = f"{cls}_{i:03d}"
            im.save(out / split / f"{name}.jpg", quality=92)
            (out / split / f"{name}.txt").write_text(cap + "\n")
            rows.append({"split": split, "class": cls, "name": name, "w": im.width, "h": im.height,
                         "bucket": bucket, "caption": cap})
            print(f"  {split}/{name} {im.width}x{im.height} | {cap[:80]}")

    with open(out / "index.csv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=["split", "class", "name", "w", "h", "bucket", "caption"])
        w.writeheader(); w.writerows(rows)
    per = {c: sum(1 for r in rows if r["class"] == c and r["split"] == "train") for c in CLASSES}
    print(f"\nитого: {sum(per.values())} train {per} + {sum(1 for r in rows if r['split']=='eval')} eval"
          f" -> {out}")


if __name__ == "__main__":
    main()

	"""Сборка набора по классам (woman/man/cat) из ~/dataset/allphoto_train — тест анатомии.

	Зачем: бабочки однородны и анатомию не показывают. Нужен набор, где модель обязана выучить
	тело: разные люди, позы, ракурсы, и капшены, которые называют, кто это и что делает.

	Классы: woman / man / cat — по N уникальных капшенов на класс (дедуп по первым 60 символам,
	иначе один photoshoot даёт десятки почти одинаковых кадров).

	python3 src/make_dataset_people.py --n 60 --out ~/dataset/testset_anatomy
	"""
	from __future__ import annotations

	import argparse
	import csv
	import math
	import random
	import re
	import shutil
	from pathlib import Path

	from PIL import Image, ImageOps

	SRC = Path.home() / "dataset" / "allphoto_train"

	CLASSES = {
	"woman": ("woman", "women", "girl"),
	"man": ("man", "men", "boy"),
	"cat": ("cat",),
	}
	# бакеты как в moss/src/make_dataset.py: обе стороны кратны 16 (латент /8, потом patch-2),
	# соотношение выбирается ближайшим по \|log(ratio_источника / ratio_бакета)\|, затем ЦЕНТР-КРОП —
	# растягивать нельзя, иначе лица и тела деформируются
	def buckets(long_side: int = 512, multiple: int = 64) -> dict:
	"""Бакеты: длинная сторона long_side, обе стороны кратны 16 (латент /8, затем /2)."""
	L = long_side
	def r16(x): return max(multiple, int(round(x / multiple)) * multiple)
	return {
	"1:1": (L, L),
	"4:3": (L, r16(L * 3 / 4)),
	"3:4": (r16(L * 3 / 4), L),
	"3:2": (L, r16(L * 2 / 3)),
	"2:3": (r16(L * 2 / 3), L),
	"16:9": (L, r16(L * 9 / 16)),
	"9:16": (r16(L * 9 / 16), L),
	}
	BUCKETS = buckets(512)

	# для людей требуем, чтобы в кадре было тело, а не только лицо: бабочки-портреты анатомию не учат
	BODY_WORDS = ("standing", "walking", "sitting", "seated", "jeans", "shorts", "dress", "skirt",
	"shirt", "legs", "arms", "hands", "wearing", "full body", "walks", "stands")


	def classify(cap: str) -> str \| None:
	low = cap.lower()
	for cls, words in CLASSES.items():
	if any(re.search(rf"\b{w}\b", low) for w in words):
	# «man» внутри «woman» уже отсечён \b; но «woman» не должна становиться «man»
	if cls == "man" and re.search(r"\bwoman\b", low):
	return "woman"
	return cls
	return None


	def trim_black(im: Image.Image, thr: float = 12.0, var: float = 8.0) -> Image.Image:
	"""Срезает однородные чёрные поля по краям (в allphoto_train их полно: 71 из 171 кадра)."""
	import numpy as np
	a = np.asarray(im.convert("L"), dtype=np.float32)
	def uniform_edge(line): # чёрная и без структуры -> рамка, а не ночная сцена
	return line.mean() < thr and line.std() < var
	top, bot, left, right = 0, a.shape[0], 0, a.shape[1]
	while top < bot - 8 and uniform_edge(a[top]): top += 1
	while bot > top + 8 and uniform_edge(a[bot - 1]): bot -= 1
	while left < right - 8 and uniform_edge(a[top:bot, left]): left += 1
	while right > left + 8 and uniform_edge(a[top:bot, right - 1]): right -= 1
	return im.crop((left, top, right, bot)) if (left or top or right != a.shape[1] or bot != a.shape[0]) else im


	def plan_size(w: int, h: int, max_upscale: float = 1.4):
	"""(w,h) -> (label, tw, th, cw, ch) или None, если источника не хватает (как в moss)."""
	ratio = w / h
	label, (tw, th) = min(BUCKETS.items(), key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1]))))
	if ratio > tw / th:
	cw, ch = round(h * tw / th), h
	else:
	cw, ch = w, round(w * th / tw)
	cw, ch = min(cw, w), min(ch, h)
	if cw * max_upscale < tw or ch * max_upscale < th:
	return None
	return label, tw, th, cw, ch


	def load_bucketed(path: Path, max_upscale: float = 1.4):
	"""EXIF-поворот -> срез чёрных полей -> ближайший бакет -> центр-кроп -> ресайз."""
	im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
	im = trim_black(im)
	plan = plan_size(im.width, im.height, max_upscale)
	if plan is None:
	return None, None
	label, tw, th, cw, ch = plan
	left, top = (im.width - cw) // 2, (im.height - ch) // 2
	return im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS), label


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--n", type=int, default=60, help="картинок на класс")
	ap.add_argument("--eval-n", type=int, default=4, help="сколько на класс отложить для показа")
	ap.add_argument("--max-upscale", type=float, default=1.4)
	ap.add_argument("--long-side", type=int, default=512)
	ap.add_argument("--multiple", type=int, default=64)
	ap.add_argument("--seed", type=int, default=0)
	ap.add_argument("--out", default=str(Path.home() / "dataset" / "testset_anatomy"))
	args = ap.parse_args()

	global BUCKETS
	BUCKETS = buckets(args.long_side, args.multiple)
	out = Path(args.out)
	for sub in ("train", "eval"):
	(out / sub).mkdir(parents=True, exist_ok=True)

	caps = sorted(SRC.glob("*.txt"))
	rng = random.Random(args.seed)
	rng.shuffle(caps)

	seen_prefix: set[str] = set()
	picked: dict[str, list[tuple[Path, str]]] = {c: [] for c in CLASSES}
	need = args.n + args.eval_n
	for f in caps:
	cap = " ".join(f.read_text().split())
	if len(cap) < 40:
	continue
	cls = classify(cap)
	if cls is None or len(picked[cls]) >= need:
	continue
	low = cap.lower()
	if cls == "cat" and re.search(r"\b(dog\|puppy\|kitten\|rabbit)\b", low) and "cat" not in low:
	continue
	if cls in ("woman", "man") and not any(w in low for w in BODY_WORDS):
	continue
	key = cap[:60].lower()
	if key in seen_prefix:
	continue
	img = f.with_suffix(".jpg")
	if not img.exists():
	continue
	seen_prefix.add(key)
	picked[cls].append((img, cap))
	if all(len(v) >= need for v in picked.values()):
	break

	rows = []
	for cls, items in picked.items():
	for i, (img, cap) in enumerate(items):
	split = "eval" if i < args.eval_n else "train"
	im, bucket = load_bucketed(img)
	if im is None:
	continue
	name = f"{cls}_{i:03d}"
	im.save(out / split / f"{name}.jpg", quality=92)
	(out / split / f"{name}.txt").write_text(cap + "\n")
	rows.append({"split": split, "class": cls, "name": name, "w": im.width, "h": im.height,
	"bucket": bucket, "caption": cap})
	print(f" {split}/{name} {im.width}x{im.height} \| {cap[:80]}")

	with open(out / "index.csv", "w", newline="") as fh:
	w = csv.DictWriter(fh, fieldnames=["split", "class", "name", "w", "h", "bucket", "caption"])
	w.writeheader(); w.writerows(rows)
	per = {c: sum(1 for r in rows if r["class"] == c and r["split"] == "train") for c in CLASSES}
	print(f"\nитого: {sum(per.values())} train {per} + {sum(1 for r in rows if r['split']=='eval')} eval"
	f" -> {out}")


	if __name__ == "__main__":
	main()