| """Сборка набора из папки «картинка + .txt»: бакеты (кратны --multiple) и центр-кроп. |
| |
| python3 src/make_dataset.py --src ~/dataset/butterflies/butterfly \ |
| --out ~/dataset/butterflies_256 --long-side 256 --eval-n 30 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import math |
| import random |
| from pathlib import Path |
|
|
| from PIL import Image, ImageOps |
|
|
| BUCKETS = {} |
|
|
|
|
| def buckets(long_side: int, multiple: int = 16) -> dict: |
| 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), |
| } |
|
|
|
|
| def plan_size(w, h, max_upscale=1.4): |
| 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 main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--src", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--long-side", type=int, default=256) |
| ap.add_argument("--eval-n", type=int, default=30) |
| ap.add_argument("--max-upscale", type=float, default=1.4) |
| ap.add_argument("--multiple", type=int, default=64, |
| help="кратность сторон: три даунсэмпла тела требуют латент, кратный 8") |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
|
|
| global BUCKETS |
| BUCKETS = buckets(args.long_side, args.multiple) |
| src, out = Path(args.src), Path(args.out) |
| for sub in ("train", "eval"): |
| (out / sub).mkdir(parents=True, exist_ok=True) |
|
|
| pairs = [] |
| for img in sorted(list(src.glob("*.png")) + list(src.glob("*.jpg"))): |
| txt = img.with_suffix(".txt") |
| if txt.exists(): |
| pairs.append((img, txt)) |
| random.Random(args.seed).shuffle(pairs) |
| print(f"пар: {len(pairs)}") |
|
|
| rows, kept = [], 0 |
| for i, (img, txt) in enumerate(pairs): |
| im = ImageOps.exif_transpose(Image.open(img)).convert("RGB") |
| plan = plan_size(im.width, im.height, args.max_upscale) |
| if plan is None: |
| continue |
| label, tw, th, cw, ch = plan |
| left, top = (im.width - cw) // 2, (im.height - ch) // 2 |
| im = im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS) |
| split = "eval" if i < args.eval_n else "train" |
| name = f"{img.stem}" |
| im.save(out / split / f"{name}.jpg", quality=95) |
| (out / split / f"{name}.txt").write_text(txt.read_text().strip() + "\n") |
| rows.append({"split": split, "name": name, "w": tw, "h": th, "bucket": label}) |
| kept += 1 |
| with open(out / "index.csv", "w", newline="") as fh: |
| w = csv.DictWriter(fh, fieldnames=["split", "name", "w", "h", "bucket"]) |
| w.writeheader(); w.writerows(rows) |
| print(f"готово: {kept} (train {kept-args.eval_n}, eval {args.eval_n}) -> {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
