src/mcss/cli.py
"""Command-line entry points for training, evaluation, data, and diagnostics."""
from __future__ import annotations
import argparse
import json
from collections.abc import Sequence
from dataclasses import replace
from pathlib import Path
from mcss.config import ExperimentConfig, load_config
from mcss.data.download import (
build_replica_plan,
download_hypersim_members,
execute_plan,
extract_replica_parts,
)
from mcss.data.hypersim import prepare_hypersim
from mcss.data.replica import prepare_replica
from mcss.diagnostics import no_bypass_audit
from mcss.engine import Trainer
def main(argv: Sequence[str] | None = None) -> int:
parser = _build_parser()
args = parser.parse_args(argv)
if args.command == "train":
trainer = Trainer(load_config(args.config))
checkpoint = trainer.train()
print(json.dumps({"checkpoint": str(checkpoint), "step": trainer.step}))
return 0
if args.command == "evaluate":
config = _with_evaluation_overrides(
load_config(args.config),
output_dir=args.output_dir,
n_samples=args.n_samples,
ray_chunk_size=args.ray_chunk_size,
)
trainer = Trainer(config)
metrics = trainer.evaluate(args.checkpoint)
print(json.dumps(metrics, sort_keys=True))
return 0
if args.command == "diagnose":
trainer = Trainer(load_config(args.config))
if args.checkpoint:
trainer.load_checkpoint(args.checkpoint, load_optimizer=False)
report = no_bypass_audit(trainer.model)
payload = {
"passed": report.passed,
"renderer_parameter_count": report.renderer_parameter_count,
"forbidden_forward_arguments": list(report.forbidden_forward_arguments),
"reasons": list(report.reasons),
}
trainer.output_dir.mkdir(parents=True, exist_ok=True)
(trainer.output_dir / "diagnostic.json").write_text(
json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
)
print(json.dumps(payload, sort_keys=True))
return 0
if args.command == "download":
return _download_command(args)
if args.command == "prepare":
return _prepare_command(args)
raise RuntimeError(f"unsupported command {args.command}")
def _build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(prog="mcss")
subparsers = parser.add_subparsers(dest="command", required=True)
for name in ("train", "evaluate", "diagnose"):
command = subparsers.add_parser(name)
command.add_argument("--config", required=True, type=Path)
if name in {"evaluate", "diagnose"}:
command.add_argument("--checkpoint", type=Path)
if name == "evaluate":
command.add_argument("--output-dir", type=Path)
command.add_argument("--n-samples", type=int)
command.add_argument("--ray-chunk-size", type=int)
download = subparsers.add_parser("download")
download.add_argument("dataset", choices=("replica", "hypersim"))
download.add_argument("--root", required=True, type=Path)
download.add_argument("--scenes", nargs="*", default=[])
download.add_argument("--camera", default="cam_00")
download.add_argument("--frames", nargs="*", type=int)
download.add_argument("--dry-run", action="store_true")
download.add_argument("--extract-to", type=Path)
download.add_argument("--remove-parts", action="store_true")
prepare = subparsers.add_parser("prepare")
prepare.add_argument("dataset", choices=("replica", "hypersim"))
prepare.add_argument("--raw-root", required=True, type=Path)
prepare.add_argument("--output-root", required=True, type=Path)
prepare.add_argument("--scenes", nargs="*")
prepare.add_argument("--camera", default="cam_00")
prepare.add_argument("--image-height", type=int)
prepare.add_argument("--image-width", type=int)
prepare.add_argument("--frame-limit", type=int)
prepare.add_argument("--frame-stride", type=int, default=1)
prepare.add_argument("--depth-scale", type=float)
return parser
def _download_command(args: argparse.Namespace) -> int:
if args.dataset == "replica":
plan = build_replica_plan(args.root)
if args.dry_run:
print(
json.dumps(
{"dataset": plan.dataset, "items": [item.url for item in plan.items]}, indent=2
)
)
return 0
results = execute_plan(plan)
if args.extract_to:
extract_replica_parts(plan, args.extract_to, remove_parts=args.remove_parts)
print(json.dumps({"dataset": "replica", "files": len(results)}))
return 0
if not args.scenes:
raise ValueError("Hypersim download requires --scenes ai_DDD_DDD [...]")
if args.dry_run:
print(
json.dumps(
{
"dataset": "hypersim",
"scenes": args.scenes,
"camera": args.camera,
"frames": args.frames,
},
indent=2,
)
)
return 0
frame_ids = None if args.frames is None else set(args.frames)
files = []
for scene in args.scenes:
files.extend(
download_hypersim_members(scene, args.root, camera=args.camera, frame_ids=frame_ids)
)
print(json.dumps({"dataset": "hypersim", "files": len(files)}))
return 0
def _prepare_command(args: argparse.Namespace) -> int:
image_size = _optional_image_size(args.image_height, args.image_width)
if args.dataset == "replica":
path = prepare_replica(
args.raw_root,
args.output_root,
frame_limit=args.frame_limit,
image_size=image_size,
depth_scale=args.depth_scale,
)
print(json.dumps({"manifest": str(path)}))
return 0
paths = prepare_hypersim(
args.raw_root,
args.output_root,
scenes=args.scenes,
camera=args.camera,
image_size=image_size,
frame_stride=args.frame_stride,
frame_limit=args.frame_limit,
)
print(json.dumps({"manifests": [str(path) for path in paths]}))
return 0
def _optional_image_size(height: int | None, width: int | None) -> tuple[int, int] | None:
if height is None and width is None:
return None
if height is None or width is None or height < 2 or width < 2:
raise ValueError("--image-height and --image-width must both be at least two")
return height, width
def _with_evaluation_overrides(
config: ExperimentConfig,
*,
output_dir: Path | None,
n_samples: int | None,
ray_chunk_size: int | None,
) -> ExperimentConfig:
if n_samples is not None and n_samples < 2:
raise ValueError("--n-samples must be at least two")
if ray_chunk_size is not None and ray_chunk_size < 1:
raise ValueError("--ray-chunk-size must be positive")
model = replace(
config.model,
n_samples=config.model.n_samples if n_samples is None else n_samples,
ray_chunk_size=(
config.model.ray_chunk_size if ray_chunk_size is None else ray_chunk_size
),
)
if output_dir is None:
return replace(config, model=model)
resolved = output_dir.resolve()
for name in ("evaluation.json", "evaluation_report.json"):
report_path = resolved / name
if report_path.exists():
raise FileExistsError(
f"refusing to overwrite existing evaluation report: {report_path}"
)
return replace(
config,
model=model,
training=replace(config.training, output_dir=str(resolved)),
)
if __name__ == "__main__":
raise SystemExit(main())