Files
StarPilot/scripts/speed_limit_vision/init_workspace.py
firestar5683 fe4f42a616 friar carl
2026-03-31 13:27:22 -05:00

96 lines
3.4 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import argparse
from pathlib import Path
if __package__ in (None, ""):
import sys
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import ( # type: ignore
BOOKMARK_MANIFEST_FIELDS,
BOOKMARK_LEADIN_MANIFEST_FIELDS,
DEFAULT_SPEED_VALUES,
VALUE_LABEL_FIELDS,
DEFAULT_WORKSPACE,
PUBLIC_CLASSIFIER_SAMPLE_FIELDS,
PUBLIC_DETECTOR_SAMPLE_FIELDS,
RAW_SOURCE_FIELDS,
detector_dataset_yaml,
ensure_dir,
resolve_workspace,
workspace_readme,
write_csv_header,
write_text,
)
else:
from .common import (
BOOKMARK_MANIFEST_FIELDS,
BOOKMARK_LEADIN_MANIFEST_FIELDS,
DEFAULT_SPEED_VALUES,
VALUE_LABEL_FIELDS,
DEFAULT_WORKSPACE,
PUBLIC_CLASSIFIER_SAMPLE_FIELDS,
PUBLIC_DETECTOR_SAMPLE_FIELDS,
RAW_SOURCE_FIELDS,
detector_dataset_yaml,
ensure_dir,
resolve_workspace,
workspace_readme,
write_csv_header,
write_text,
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Initialize a speed-limit detector/classifier training workspace.")
parser.add_argument("--workspace", type=Path, default=DEFAULT_WORKSPACE, help="Workspace root. Defaults to .tmp/speed_limit_training under the repo root.")
parser.add_argument("--force", action="store_true", help="Overwrite generated template files if they already exist.")
parser.add_argument("--speed-values", nargs="+", type=int, default=list(DEFAULT_SPEED_VALUES), help="Classifier speed values to document in the workspace README.")
return parser.parse_args()
def main() -> int:
args = parse_args()
workspace = resolve_workspace(args.workspace)
for relative_dir in (
"detector/images/train",
"detector/images/val",
"detector/labels/train",
"detector/labels/val",
"classifier/train",
"classifier/val",
"review/images",
"review/leadins/frames",
"review/leadins/contact_sheets",
"manifests",
"raw",
"staging",
"exports",
"runs",
):
ensure_dir(workspace / relative_dir)
write_text(workspace / "detector" / "dataset.yaml", detector_dataset_yaml(workspace), force=args.force)
write_text(workspace / "README.md", workspace_readme(tuple(args.speed_values)), force=args.force)
write_csv_header(workspace / "review" / "bookmarks.csv", BOOKMARK_MANIFEST_FIELDS, force=args.force)
write_csv_header(workspace / "review" / "bookmark_leadins.csv", BOOKMARK_LEADIN_MANIFEST_FIELDS, force=args.force)
write_csv_header(workspace / "classifier" / "value_labels.csv", VALUE_LABEL_FIELDS, force=args.force)
write_csv_header(workspace / "manifests" / "raw_sources.csv", RAW_SOURCE_FIELDS, force=args.force)
write_csv_header(workspace / "manifests" / "public_detector_samples.csv", PUBLIC_DETECTOR_SAMPLE_FIELDS, force=args.force)
write_csv_header(workspace / "manifests" / "public_classifier_samples.csv", PUBLIC_CLASSIFIER_SAMPLE_FIELDS, force=args.force)
print(f"Initialized speed-limit training workspace at {workspace}")
print(f" detector dataset: {workspace / 'detector/dataset.yaml'}")
print(f" review manifest: {workspace / 'review/bookmarks.csv'}")
print(f" lead-in manifest: {workspace / 'review/bookmark_leadins.csv'}")
print(f" value labels: {workspace / 'classifier/value_labels.csv'}")
print(f" raw sources: {workspace / 'manifests/raw_sources.csv'}")
return 0
if __name__ == "__main__":
raise SystemExit(main())