scripts/init_local_focus.py

"""
Create local dataset folders for focused plastic/glass improvement work.

Usage:
    python scripts/init_local_focus.py
    python scripts/init_local_focus.py --root data
"""

import argparse
from pathlib import Path


LOCAL_BOOST_CLASSES = ["plastic", "glass"]
LOCAL_TEST_CLASSES = ["plastic", "glass", "unknown"]


def ensure_dirs(paths: list[Path]) -> None:
    for path in paths:
        path.mkdir(parents=True, exist_ok=True)
        keep_path = path / ".gitkeep"
        keep_path.touch(exist_ok=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Create focused local dataset folders.")
    parser.add_argument("--root", default="data", help="Base data directory")
    args = parser.parse_args()

    root = Path(args.root)
    local_boost_root = root / "local_boost"
    local_test_root = root / "local_test_focus"

    ensure_dirs([local_boost_root / class_name for class_name in LOCAL_BOOST_CLASSES])
    ensure_dirs([local_test_root / class_name for class_name in LOCAL_TEST_CLASSES])

    print("Created local dataset folders:")
    print(f"  boost: {local_boost_root.resolve()}")
    print(f"  test : {local_test_root.resolve()}")
    print("\nRecommended use:")
    print("  data/local_boost/plastic -> hard plastic examples for retraining")
    print("  data/local_boost/glass   -> hard glass examples for retraining")
    print("  data/local_test_focus/*  -> held-out evaluation images, not training")


if __name__ == "__main__":
    main()