80 lines
2.2 KiB
Python
80 lines
2.2 KiB
Python
"""
|
|
generate_splits_txt.py
|
|
|
|
Generates a splits.txt file listing all label filenames across train/test/val
|
|
splits in the format:
|
|
|
|
split/labels/<filename>.txt
|
|
|
|
One entry per line, ordered train → test → val. This file can be consumed by a
|
|
shell script to recreate the directory structure and create symlinks, e.g.:
|
|
|
|
while IFS= read -r entry; do
|
|
split=$(echo "$entry" | cut -d/ -f1)
|
|
fname=$(basename "$entry" .txt)
|
|
ln -s "$SOURCE/labels/$fname.txt" "splits/$split/labels/$fname.txt"
|
|
ln -s "$SOURCE/images/$fname.jpg" "splits/$split/images/$fname.jpg"
|
|
done < splits.txt
|
|
"""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
|
|
SPLITS = ["train", "test", "val"]
|
|
|
|
|
|
def generate_splits_txt(splits_dir: Path, output_file: Path) -> None:
|
|
entries = []
|
|
counts = {}
|
|
|
|
for split in SPLITS:
|
|
labels_dir = splits_dir / split / "labels"
|
|
if not labels_dir.is_dir():
|
|
print(f"Warning: {labels_dir} not found, skipping.")
|
|
counts[split] = 0
|
|
continue
|
|
|
|
files = sorted(labels_dir.glob("*.txt"))
|
|
for f in files:
|
|
entries.append(f"{split}/labels/{f.name}")
|
|
counts[split] = len(files)
|
|
|
|
output_file.write_text("\n".join(entries) + "\n")
|
|
|
|
total = sum(counts.values())
|
|
print(f"Written {total} entries to {output_file}")
|
|
for split in SPLITS:
|
|
print(f" {split}: {counts.get(split, 0)}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Generate splits.txt from a splits directory.")
|
|
parser.add_argument(
|
|
"--input_dir",
|
|
type=Path,
|
|
help="Path to the input folder containing train/val/test splits",
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
type=Path,
|
|
default=Path(__file__).resolve().parents[2] / "data" / "splits",
|
|
help="Path to the output folder (default: <project_root>/data/splits)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
splits_dir = args.input_dir
|
|
output_dir = args.output_dir
|
|
|
|
if not splits_dir.is_dir():
|
|
raise SystemExit(f"Error: input folder not found: {splits_dir}")
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_file = output_dir / f"{splits_dir.name}_split.txt"
|
|
|
|
generate_splits_txt(splits_dir, output_file)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|