Spaces:
Running on Zero
Running on Zero
| 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() | |