justrun.py

import os, random
from pathlib import Path
from collections import defaultdict

random.seed(42)

# ----- Edit caps here -----
CAP_IMG_PER_CLASS = 40000
CAP_VID_PER_CLASS = 800

# ----- Your sources -----
IMG_AI = [
  "/rm_eaby/All Datasets/DS_BigGAN/imagenet_ai_0419_biggan/train/ai",
  "/rm_eaby/All Datasets/DS_MidJourney/imagenet_midjourney/train/ai",
  "/rm_eaby/All Datasets/DS_VQDM/imagenet_ai_0419_vqdm/train/ai",
]
IMG_REAL = [
  "/rm_eaby/All Datasets/DS_BigGAN/imagenet_ai_0419_biggan/train/nature",
  "/rm_eaby/All Datasets/DS_MidJourney/imagenet_midjourney/train/nature",
  "/rm_eaby/All Datasets/DS_VQDM/imagenet_ai_0419_vqdm/train/nature",
]
VID_AI = [
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/Celeb-DF-v2/Celeb-synthesis/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/BDAnimateDiffLightning_Generated/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/CogVideoX5B_Generated/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/RunwayML_Generated/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/StableDiffusion_Generated/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/Veo_Generated/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/VideoPoet_Generated/",
]
VID_REAL = [
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/Celeb-DF-v2/Celeb-real/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/Celeb-DF-v2/YouTube-real/",
  "/rm_eaby/DeepFakeVideoClassifier_Project/Video Classification datasets/deepaction_v1 datasets/deepaction_v1/Pexels_realVideos/",
]

IMAGE_EXTS = {".jpg",".jpeg",".png",".bmp",".tiff",".webp",".gif",".jpeg"," .JPEG",".JPG",".PNG",".BMP",".WEBP",".TIFF",".GIF",".JPEG"}
VIDEO_EXTS = {".mp4",".avi",".mov",".mkv",".webm",".mpg",".mpeg",".m4v",".MP4",".AVI",".MOV",".MKV",".WEBM",".MPG",".MPEG",".M4V"}

def list_media(root, is_video=False):
    root = Path(root)
    if not root.exists():
        return []
    out = []
    for p in root.rglob("*"):
        if p.is_file():
            ext = p.suffix
            if (ext.lower() in VIDEO_EXTS) if is_video else (ext.lower() in IMAGE_EXTS):
                try:
                    out.append(str(p.resolve()))
                except Exception:
                    out.append(str(p))
    return out

def bucketize(paths, label, source_tag):
    # returns list of tuples (path, label, source)
    return [(p, label, source_tag) for p in paths]

def gather_group(sources, label, is_video):
    items = []
    for s in sources:
        tag = Path(s.rstrip("/")).name
        paths = list_media(s, is_video=is_video)
        # unique by full path within this source
        paths = sorted(set(paths))
        items += bucketize(paths, label, tag)
    return items

# Gather
img_ai   = gather_group(IMG_AI,   1, is_video=False)
img_real = gather_group(IMG_REAL, 0, is_video=False)
vid_ai   = gather_group(VID_AI,   1, is_video=True)
vid_real = gather_group(VID_REAL, 0, is_video=True)

# De-dup across all sources by path
def dedup(items):
    seen = set(); out = []
    for row in items:
        if row[0] in seen: 
            continue
        seen.add(row[0]); out.append(row)
    return out

img_ai   = dedup(img_ai)
img_real = dedup(img_real)
vid_ai   = dedup(vid_ai)
vid_real = dedup(vid_real)

print(f"[scan] images AI={len(img_ai)}  REAL={len(img_real)}")
print(f"[scan] videos AI={len(vid_ai)}  REAL={len(vid_real)}")

# Cap per class (uniform mixture across sources)
def cap_by_source(items, cap):
    by_source = defaultdict(list)
    for path, lbl, src in items:
        by_source[src].append((path,lbl,src))
    # compute proportional quotas
    total = sum(len(v) for v in by_source.values())
    if total == 0: return []
    # initial alloc
    alloc = {s: int(round(len(v) / total * cap)) for s,v in by_source.items()}
    # adjust rounding to hit cap
    diff = cap - sum(alloc.values())
    # distribute remainder
    for s in sorted(by_source.keys()):
        if diff == 0: break
        alloc[s] += 1 if diff > 0 else -1
        diff += -1 if diff > 0 else 1
    out = []
    for s, rows in by_source.items():
        random.shuffle(rows)
        out += rows[:max(0, alloc[s])]
    return out[:cap]

img_ai_cap   = cap_by_source(img_ai,   CAP_IMG_PER_CLASS)
img_real_cap = cap_by_source(img_real, CAP_IMG_PER_CLASS)
vid_ai_cap   = cap_by_source(vid_ai,   CAP_VID_PER_CLASS)
vid_real_cap = cap_by_source(vid_real, CAP_VID_PER_CLASS)

print(f"[cap] images AI={len(img_ai_cap)} REAL={len(img_real_cap)}  (target {CAP_IMG_PER_CLASS})")
print(f"[cap] videos AI={len(vid_ai_cap)} REAL={len(vid_real_cap)}  (target {CAP_VID_PER_CLASS})")

# Combine
all_items = img_ai_cap + img_real_cap + vid_ai_cap + vid_real_cap
random.shuffle(all_items)

# Split 80/10/10 preserving label & modality by grouping on (label, is_video)
def modality(path):
    ext = Path(path).suffix.lower()
    return 1 if ext in VIDEO_EXTS else 0

groups = defaultdict(list)
for path,lbl,src in all_items:
    groups[(lbl, modality(path))].append((path,lbl,src))

train, val, test = [], [], []
for key, rows in groups.items():
    n = len(rows)
    n_train = int(0.8 * n)
    n_val   = int(0.1 * n)
    n_test  = n - n_train - n_val
    train += rows[:n_train]
    val   += rows[n_train:n_train+n_val]
    test  += rows[n_train+n_val:]

def write_tsv(rows, out_path):
    outp = Path(out_path)
    outp.parent.mkdir(parents=True, exist_ok=True)
    with open(outp, "w") as f:
        for p,l,s in rows:
            f.write(f"{p}\t{l}\t{s}\n")

out_dir = Path("runs/filelists_balanced40k_fixed")
write_tsv(train, out_dir/"train.tsv")
write_tsv(val,   out_dir/"val.tsv")
write_tsv(test,  out_dir/"test.tsv")

print(f"[write] train={len(train)} val={len(val)} test={len(test)} → {out_dir}")