Spaces:
Sleeping
Sleeping
File size: 2,203 Bytes
127b976 | 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 |
from __future__ import annotations
import argparse
from pathlib import Path
import pandas as pd
REQUIRED_COLUMNS = {"file_path", "binary_label", "attack_type", "start_fake", "end_fake", "dataset", "split"}
def print_counts(df, title):
print("\n" + "=" * 70)
print(title)
print("=" * 70)
print("Rows:", len(df))
print("\nBinary label counts:")
print(df["binary_label"].value_counts().sort_index())
print("\nDataset counts:")
print(df["dataset"].value_counts())
print("\nDataset x label counts:")
print(pd.crosstab(df["dataset"], df["binary_label"]))
def main():
p = argparse.ArgumentParser(description="Balance metadata by keeping all real rows and downsampling fake rows.")
p.add_argument("--input", required=True)
p.add_argument("--out", required=True)
p.add_argument("--fake-multiplier", type=float, default=1.0, help="1.0 means fake count roughly equals real count")
p.add_argument("--cap-wavefake", type=int, default=0)
p.add_argument("--seed", type=int, default=42)
args = p.parse_args()
df = pd.read_csv(args.input, low_memory=False)
df["binary_label"] = df["binary_label"].astype(int)
print_counts(df, "Before balancing")
if args.cap_wavefake > 0:
mask = df["dataset"].astype(str).str.lower() == "wavefake"
wave = df[mask]
other = df[~mask]
if len(wave) > args.cap_wavefake:
wave = wave.sample(n=args.cap_wavefake, random_state=args.seed)
df = pd.concat([other, wave], ignore_index=True)
print_counts(df, f"After WaveFake cap: {args.cap_wavefake}")
real = df[df["binary_label"] == 0]
fake = df[df["binary_label"] == 1]
max_fake = int(len(real) * args.fake_multiplier)
fake = fake.sample(n=min(len(fake), max_fake), random_state=args.seed)
out = pd.concat([real, fake], ignore_index=True).sample(frac=1.0, random_state=args.seed).reset_index(drop=True)
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
out.to_csv(args.out, index=False)
print_counts(out, "After balancing")
print("Saved:", args.out)
if __name__ == "__main__":
main()
|