File size: 2,297 Bytes
e23cf4c | 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 | """Tạo split CÂN BẰNG 1:1 từ các CSV split đầy đủ (data/splits/*.csv).
Cân bằng bằng cách chỉ giữ:
- real = wiki (toàn bộ ảnh thật)
- fake = 1 loại duy nhất (mặc định: insight — swap face)
Identity vẫn tách biệt giữa train/val/test (kế thừa từ split gốc theo nhân vật).
Output: data/splits/{train,val,test}_{fake_type}.csv
"""
import argparse
import csv
import os
import re
EXTENSIONS = (".jpg", ".jpeg", ".png")
REAL_CAT = "wiki"
def keep_row(path: str, fake_type: str) -> bool:
"""Giữ ảnh nếu thuộc real (wiki) hoặc đúng loại fake được chọn."""
return f"/hug/{REAL_CAT}/" in path or f"/hug/{fake_type}/" in path
def main():
parser = argparse.ArgumentParser(description="Tạo split cân bằng 1:1")
parser.add_argument("--splits-dir", default="data/splits")
parser.add_argument("--fake-type", default="insight", choices=["insight", "inpainting", "text2img"])
parser.add_argument("--out-dir", default="data/splits")
args = parser.parse_args()
for name in ("train", "val", "test"):
src = os.path.join(args.splits_dir, f"{name}.csv")
out = os.path.join(args.out_dir, f"{name}_{args.fake_type}.csv")
kept, dropped = [], 0
with open(src, newline="") as f:
reader = csv.reader(f)
header = next(reader, None)
for row in reader:
if len(row) < 2:
continue
path, label = row[0], int(row[1])
if not path.lower().endswith(EXTENSIONS):
continue
if keep_row(path, args.fake_type):
kept.append((path, label))
else:
dropped += 1
with open(out, "w", newline="") as f:
w = csv.writer(f)
w.writerow(["path", "label"])
w.writerows(kept)
n_real = sum(1 for _, lbl in kept if lbl == 0)
n_fake = sum(1 for _, lbl in kept if lbl == 1)
print(f" {name:>5}: {len(kept):>6} ảnh (giữ) | real={n_real} fake={n_fake} | bỏ {dropped} loại fake khác")
print(f"\nĐã tạo split cân bằng (real=wiki, fake={args.fake_type}) vào {args.out_dir}/")
if __name__ == "__main__":
main()
|