File size: 3,413 Bytes
b4f5f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """Сборка набора из папки «картинка + .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()
|