File size: 4,187 Bytes
8e932ab | 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 | """Create a stratified train / val / test split manifest for the Augmented Dataset.
Outputs a single JSON manifest so the same split is reused by every model run
(training, k-fold CV, and the final independent-test evaluation).
The independent test set is held out FIRST and is never used during k-fold CV.
The k-fold CV runs on the remaining train+val pool (the script also stores
five stratified train/val folds so they can be reproduced exactly).
"""
from __future__ import annotations
import argparse
import json
import random
from collections import Counter, defaultdict
from pathlib import Path
from sklearn.model_selection import StratifiedKFold, train_test_split
def collect_samples(data_dir: Path) -> tuple[list[tuple[str, str]], list[str]]:
classes = sorted([p.name for p in data_dir.iterdir() if p.is_dir()])
samples: list[tuple[str, str]] = []
valid_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}
for cls in classes:
for image_path in sorted((data_dir / cls).iterdir()):
if image_path.suffix.lower() in valid_exts:
samples.append((str(image_path.relative_to(data_dir)), cls))
return samples, classes
def stratified_split(samples, test_size, val_size, seed):
paths = [s[0] for s in samples]
labels = [s[1] for s in samples]
paths_pool, paths_test, labels_pool, labels_test = train_test_split(
paths, labels, test_size=test_size, stratify=labels, random_state=seed
)
relative_val = val_size / (1.0 - test_size)
paths_train, paths_val, labels_train, labels_val = train_test_split(
paths_pool, labels_pool, test_size=relative_val, stratify=labels_pool, random_state=seed
)
return (paths_train, labels_train), (paths_val, labels_val), (paths_test, labels_test)
def kfold_indices(paths, labels, folds, seed):
skf = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed)
out = []
for k, (train_idx, val_idx) in enumerate(skf.split(paths, labels), start=1):
out.append({
"fold": k,
"train": [int(i) for i in train_idx],
"val": [int(i) for i in val_idx],
})
return out
def class_distribution(labels):
return dict(Counter(labels))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--data-dir", default="Database/Augmented_Dataset")
parser.add_argument("--output", default="holdout_split.json")
parser.add_argument("--test-size", type=float, default=0.15)
parser.add_argument("--val-size", type=float, default=0.15)
parser.add_argument("--folds", type=int, default=5)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
data_dir = Path(args.data_dir).resolve()
samples, classes = collect_samples(data_dir)
print(f"Dataset: {data_dir}")
print(f" total images: {len(samples)}")
print(f" classes: {classes}")
(train, val, test) = stratified_split(samples, args.test_size, args.val_size, args.seed)
folds = kfold_indices(
train[0] + val[0],
train[1] + val[1],
args.folds,
args.seed,
)
manifest = {
"data_dir": str(data_dir),
"classes": classes,
"seed": args.seed,
"test_size": args.test_size,
"val_size": args.val_size,
"splits": {
"train": list(zip(train[0], train[1])),
"val": list(zip(val[0], val[1])),
"test": list(zip(test[0], test[1])),
},
"kfold": {
"folds": args.folds,
"pool_paths": train[0] + val[0],
"pool_labels": train[1] + val[1],
"indices": folds,
},
"class_distribution": {
"train": class_distribution(train[1]),
"val": class_distribution(val[1]),
"test": class_distribution(test[1]),
},
}
out_path = Path(args.output).resolve()
out_path.write_text(json.dumps(manifest, indent=2))
print(f"Manifest written: {out_path}")
for split_name in ("train", "val", "test"):
print(f" {split_name}: {len(manifest['splits'][split_name])} images")
if __name__ == "__main__":
main()
|