File size: 5,070 Bytes
a2ffd07 | 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """
Upload the bathroom-toilet dataset with train/val split to HuggingFace.
Replicates the exact split used in training (see experiment/data/datasets.py):
sklearn.model_selection.train_test_split(image_ids, test_size=0.2, random_state=42)
Usage:
python -m experiment.scripts.upload_split_to_hf \
--csv_path CC3M-Dataset/bathroom_filter/bathroom_toilet_labels.csv \
--image_dir CC3M-Dataset/cc3m_images/train \
--repo_id <your-hf-username>/bathroom-toilet-cc3m
"""
import argparse
import csv
import os
from datasets import Dataset, DatasetDict, Features, Value, Image, ClassLabel
from sklearn.model_selection import train_test_split
from huggingface_hub import login
# Must match experiment/data/datasets.py
SPLIT_SEED = 42
SPLIT_TEST_SIZE = 0.2
def load_rows(csv_path: str, image_dir: str):
"""Load all rows from the labels CSV, filtering out non-bathroom and missing images."""
rows = []
missing = 0
skipped_non_bathroom = 0
with open(csv_path, "r") as f:
reader = csv.DictReader(f)
for row in reader:
bathroom = int(row["bathroom"])
toilet = int(row["toilet"])
if bathroom == 0 and toilet == 0:
skipped_non_bathroom += 1
continue
image_path = os.path.join(image_dir, f"{row['image_id']}.jpg")
if not os.path.exists(image_path):
missing += 1
continue
rows.append({
"image_id": row["image_id"],
"image": image_path,
"bathroom": bathroom,
"toilet": toilet,
})
print(f"Loaded {len(rows)} rows "
f"({skipped_non_bathroom} non-bathroom skipped, {missing} images not found)")
return rows
def split_rows(rows: list[dict]) -> tuple[list[dict], list[dict]]:
"""Replicate the exact train/val split from datasets.py."""
all_ids = [r["image_id"] for r in rows]
train_ids, val_ids = train_test_split(
all_ids, test_size=SPLIT_TEST_SIZE, random_state=SPLIT_SEED,
)
train_set = set(train_ids)
val_set = set(val_ids)
train_rows = [r for r in rows if r["image_id"] in train_set]
val_rows = [r for r in rows if r["image_id"] in val_set]
print(f"Split: {len(train_rows)} train, {len(val_rows)} val "
f"(seed={SPLIT_SEED}, test_size={SPLIT_TEST_SIZE})")
return train_rows, val_rows
def rows_to_dataset(rows: list[dict]) -> Dataset:
"""Convert list of row dicts to a HuggingFace Dataset."""
return Dataset.from_dict(
{
"image_id": [r["image_id"] for r in rows],
"image": [r["image"] for r in rows],
"bathroom": [r["bathroom"] for r in rows],
"toilet": [r["toilet"] for r in rows],
},
features=Features({
"image_id": Value("string"),
"image": Image(),
"bathroom": ClassLabel(names=["no", "yes"]),
"toilet": ClassLabel(names=["no", "yes"]),
}),
)
def print_report(train_rows, val_rows):
"""Print category breakdown matching the training script style."""
for name, rows in [("train", train_rows), ("val", val_rows)]:
cats: dict[str, int] = {}
for r in rows:
key = f"bathroom={r['bathroom']},toilet={r['toilet']}"
cats[key] = cats.get(key, 0) + 1
print(f" {name}: {cats}")
def main():
parser = argparse.ArgumentParser(
description="Upload bathroom-toilet dataset splits to HuggingFace"
)
parser.add_argument(
"--csv_path",
type=str,
default="CC3M-Dataset/bathroom_filter/bathroom_toilet_labels.csv",
)
parser.add_argument(
"--image_dir",
type=str,
default="CC3M-Dataset/cc3m_images/train",
)
parser.add_argument(
"--repo_id",
type=str,
required=True,
help="HuggingFace repo id, e.g. your-username/bathroom-toilet-cc3m",
)
parser.add_argument(
"--private",
action="store_true",
help="Make the dataset private on HuggingFace",
)
args = parser.parse_args()
# Auth
hf_token = os.environ.get("HUGGING_FACE_API_KEY")
if hf_token:
login(token=hf_token)
else:
print("No HUGGING_FACE_API_KEY in env, using cached HF credentials")
# Load and split
rows = load_rows(args.csv_path, args.image_dir)
if len(rows) == 0:
print("ERROR: No data loaded. Check csv_path and image_dir.")
return
train_rows, val_rows = split_rows(rows)
print_report(train_rows, val_rows)
# Build HF datasets
print("\nBuilding HuggingFace datasets...")
ds = DatasetDict({
"train": rows_to_dataset(train_rows),
"val": rows_to_dataset(val_rows),
})
print(ds)
# Upload
print(f"\nPushing to {args.repo_id}...")
ds.push_to_hub(args.repo_id, private=args.private)
print(f"Done! Dataset available at: https://huggingface.co/datasets/{args.repo_id}")
if __name__ == "__main__":
main()
|