File size: 4,149 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 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 96 97 98 99 100 101 102 | """Chia dataset DeepFakeFace (data/hug) thành train / val / test THEO NHÂN VẬT.
Lý do chia theo nhân vật (folder 00–99): cùng 1 người xuất hiện ở cả 4 category
(wiki=real, insight/inpainting/text2img=fake). Nếu chia ngẫu nhiên theo ảnh, model
sẽ học gộp 1 người ở cả train lẫn test → điểm bị đánh giá cao ảo.
Tỷ lệ: 100 nhân vật → 72 train / 8 val / 20 test (test = 20%).
Output (data/splits/*.csv): cột `path,label` (label: 0=real, 1=fake), path tương đối
so với thư mục project. Không nhân đôi ảnh — chỉ là danh sách đường dẫn.
"""
import argparse
import csv
import os
import random
from collections import defaultdict
# wiki = ảnh thật (real), 3 category còn lại = deepfake (fake)
CATEGORIES = {"wiki": 0, "insight": 1, "inpainting": 1, "text2img": 1}
EXTENSIONS = (".jpg", ".jpeg", ".png")
def discover_samples(data_dir: str):
"""Gom toàn bộ ảnh theo nhân vật. Trả về dict {identity: [(path, label), ...]}."""
samples = defaultdict(list)
for cat, label in CATEGORIES.items():
cat_dir = os.path.join(data_dir, cat)
if not os.path.isdir(cat_dir):
print(f"[Cảnh báo] Không có thư mục: {cat_dir}")
continue
for ident in sorted(os.listdir(cat_dir)):
idir = os.path.join(cat_dir, ident)
if not os.path.isdir(idir):
continue
for f in sorted(os.listdir(idir)):
if f.lower().endswith(EXTENSIONS):
samples[ident].append((os.path.join(cat_dir, ident, f), label))
return samples
def write_csv(out_path: str, rows):
os.makedirs(os.path.dirname(out_path), exist_ok=True)
with open(out_path, "w", newline="") as f:
w = csv.writer(f)
w.writerow(["path", "label"])
w.writerows(rows)
def main():
parser = argparse.ArgumentParser(description="Chia dataset theo nhân vật")
parser.add_argument("--data-dir", default="data/hug")
parser.add_argument("--out-dir", default="data/splits")
parser.add_argument("--test-identities", type=int, default=20, help="Số nhân vật làm test (20 = 20%)")
parser.add_argument("--val-identities", type=int, default=8, help="Số nhân vật làm val (lấy từ phần train)")
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
samples = discover_samples(args.data_dir)
idents = sorted(samples.keys())
print(f"Nhân vật: {len(idents)} | Tổng ảnh: {sum(len(v) for v in samples.values())}")
# ---- Chia ngẫu nhiên theo nhân vật (có seed để tái lập) ----
rng = random.Random(args.seed)
shuffled = idents[:]
rng.shuffle(shuffled)
test_ids = set(shuffled[: args.test_identities])
rest = shuffled[args.test_identities:]
val_ids = set(rest[: args.val_identities])
train_ids = set(rest[args.val_identities:])
splits = {
"train": train_ids,
"val": val_ids,
"test": test_ids,
}
print(f"\nSplit: train={len(train_ids)} nhân vật | val={len(val_ids)} | test={len(test_ids)}")
# ---- Xuất CSV ----
for name, id_set in splits.items():
rows = []
for ident in sorted(id_set):
rows.extend(samples[ident])
# path tương đối so với project (bỏ tiền tố "./")
rows = [(p[2:] if p.startswith("./") else p, lbl) for p, lbl in rows]
write_csv(os.path.join(args.out_dir, f"{name}.csv"), rows)
n_real = sum(1 for _, lbl in rows if lbl == 0)
n_fake = sum(1 for _, lbl in rows if lbl == 1)
print(f" {name:>5}: {len(rows):>7} ảnh | real={n_real:>6} | fake={n_fake:>6} (tỷ lệ 1:{n_fake / n_real:.1f})")
# ---- Lưu thông tin split ----
summary = {name: sorted(s) for name, s in splits.items()}
import json
with open(os.path.join(args.out_dir, "split_info.json"), "w") as f:
json.dump({"seed": args.seed, "splits": summary}, f, indent=2)
print(f"\nĐã lưu CSV + split_info.json vào {args.out_dir}/")
if __name__ == "__main__":
main()
|